Use .parameter method in Python on a c++ nn::Module exposed through torch::jit::CustomClassHolder

Mat32: <p>Hello, I want to optimize in Python the weights of a Torch module implemented in c++. Here is the idea, I generate a dataset yrand = Wsol * xrand, implement in c++ a module that computes W * x product and I want to find the best W that minimise || W * xrand - yrand ||</p>
<p>%%%%%%%%%%%<br>
%myTorchClass.cpp%<br>
%%%%%%%%%%%</p>
<pre><code class="lang-auto">struct MyTorchClass : torch::jit::CustomClassHolder
{
    torch::Tensor g(torch::Tensor x) const
    {
        return m_W.mm(x);
    }

    void init()
    {
        m_W = torch::zeros({ 3, 3}, torch::TensorOptions().dtype(torch::kFloat64).requires_grad(true));
        m_module.register_parameter("W", m_W, true);
    }

    // Expose a method to get parameters, callable from Python
    at::Tensor&amp; parameters() const {
        const auto&amp; ret = m_module.named_parameters();
        return m_module.named_parameters()["W"];
    }

    torch::Tensor m_W;
    torch::nn::Module m_module;
};

TORCH_LIBRARY(mynamespace, m)
{
    m.class_&lt;MyTorchClass&gt;("myclassname")
        .def(torch::init&lt;&gt;())
        .def("init", &amp;MyTorchClass::init)
        .def("g", &amp;MyTorchClass::g)
        .def("parameters", MyTorchClass::parameters);
}
</code></pre>
<p>%%%%%%%<br>
% main.py%%<br>
%%%%%%%</p>
<pre><code class="lang-auto">xrand = torch.rand((3, 1000), dtype=torch.float64)
Wsol = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=torch.float64)
yrand = torch.matmul(Wsol, xrand)

class myObjective(torch.nn.Module):
    def __init__(self, inputs: torch.Tensor, targets: torch.Tensor):
        super().__init__()
        self.my_ops = my_torch_ops()
        self.inputs = inputs
        self.targets = targets
    def forward(self, x):
        return (self.my_ops.g(self.inputs, self.inputs) - self.targets).norm()

my_objective = myObjective(xrand, yrand)

optimizer = torch.optim.SGD(params=my_objective.my_ops.parameters, lr=0.01)

# Optimization loop
for i in range(1000):
    # Zero the gradients
    optimizer.zero_grad()

    # Compute the value of the objective function
    y = my_objective.forward()

    # Compute the gradients
    y.backward()

    # Update the parameters
    optimizer.step()

</code></pre>
<p>I validated that loading custom class mechanism through ‘torch.ops.load_library’ works as intended. The .forward() of my_objective works as well.<br>
But I am missing something regarding the way to manage parameters. This code crashes at</p>
<p><code>optimizer = torch.optim.SGD(params=my_objective.my_ops.parameters, lr=0.01)</code></p>
<p>with the error</p>
<pre><code class="lang-auto">  File "main.py", line 61, in &lt;module&gt;
    optimizer = torch.optim.SGD(params=my_objective.my_ops.parameters, lr=0.01)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/torch/optim/sgd.py", line 27, in __init__
    super().__init__(params, defaults)
  File ".../lib/python3.11/site-packages/torch/optim/optimizer.py", line 185, in __init__
    param_groups = list(params)
                   ^^^^^^^^^^^^
TypeError: 'torch._C.ScriptMethod' object is not iterable
</code></pre>
<p>Is that kind of design viable and how to fix this issue? Thx</p>

ptrblck: <p>Try to use <code>optimizer = torch.optim.SGD(params=my_objective.my_ops.parameters(), lr=0.01)</code> (note the call into <code>.parameters()</code>).</p>

Mat32: <p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a> for your fast reply, you’re right</p>
<pre><code class="lang-auto">optimizer = torch.optim.SGD(params=[my_objective.rnn_ops.rnn_obj.parameters()], lr=0.01)
</code></pre>
<p>fixed the issue and the optimization gives me the Wsol !</p>
<p>When I started working on that python / c++ module interaction I wasn’t sure about the overall design. What I understand is that my cpp class has to herit from jit:CustomClassHolder to be loaded in Python. It’s cool but lacks the handy ‘.parameter’ mechanism of nn::Module.<br>
So then I added a nn::Module member to the class and a parameters() method to expose the parameters() method of a nn::Module.<br>
Isn’t there a more direct way to expose a nn::Module, using multiple inheritance in MyTorchClass for example (I know there is a torch::jit::Module), or anything else ?</p>