Custom C++ extension for Optimizer. state_dict not updated

Raul_Murillo: <p>Hello all,</p>
<p>I am trying to build a <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" rel="noopener nofollow ugc">custom C++ extension</a> with PyTorch.<br>
In particular, I am trying to define a new optimizer <code>step()</code> function, so I started replicating the functionality of <a href="https://pytorch.org/docs/stable/generated/torch.optim.SGD.html" rel="noopener nofollow ugc">SGD optimizer</a> <code>step()</code> as a C++ extension.</p>
<p>Here is the C++ code:</p>
<pre><code class="lang-cpp">#include &lt;torch/extension.h&gt;
#include &lt;iostream&gt;
#include &lt;vector&gt;

void my_sgd_function(std::vector&lt;torch::Tensor&gt; &amp;params,
                  const std::vector&lt;torch::Tensor&gt; &amp;d_p_list,
                  std::vector&lt;std::optional&lt;torch::Tensor&gt;&gt; &amp;momentum_buffer_list,
                  double weight_decay,
                  double momentum,
                  double lr,
                  double dampening,
                  bool nesterov)
{
    const size_t num_params = params.size();

    for (size_t i = 0; i &lt; num_params; ++i)
    {
        torch::Tensor d_p = d_p_list[i];
        d_p.add_(params[i], weight_decay);

        if (momentum != 0)
        {
            torch::Tensor buf = torch::empty_like(d_p);

            if (!momentum_buffer_list[i].has_value())
            {
                buf = d_p.clone().detach();
                momentum_buffer_list[i] = buf;
            }
            else
            {
                buf = momentum_buffer_list[i].value();
                buf.mul_(momentum).add_(d_p, 1 - dampening);
            }

            if (nesterov)
            {
                d_p.add_(buf, momentum);
            }
            else
            {
                d_p = buf;
            }
        }

        params[i].add_(d_p, -lr);
    }
}

// Binding
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("my_sgd", &amp;my_sgd_function, "my sgd computation (CPU)");
}
</code></pre>
<p>My C function tries to mimic the behavior of <a href="https://github.com/pytorch/pytorch/blob/59ccc786dff9657ac85aae2b333d521728608cd6/torch/optim/sgd.py#L234" rel="noopener nofollow ugc">_single_tensor_sgd()</a>.<br>
After compiling and binding, here is the Python code from which I want to use the custom C++ function (extracted from <a href="https://github.com/pytorch/pytorch/blob/master/torch/optim/sgd.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch/sgd.py at master · pytorch/pytorch · GitHub</a>):</p>
<pre><code class="lang-python">import torch
import torch.optim._functional as F
from torch import Tensor
from typing import List, Optional
# My module!
import sgd_cpu

class mySGD(optim.Optimizer):
    def __init__(self, params, lr, momentum=0, dampening=0,
                 weight_decay=0, nesterov=False, maximize: bool = False):
        if lr &lt; 0.0:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if momentum &lt; 0.0:
            raise ValueError("Invalid momentum value: {}".format(momentum))
        if weight_decay &lt; 0.0:
            raise ValueError(
                "Invalid weight_decay value: {}".format(weight_decay))

        defaults = dict(lr=lr, momentum=momentum, dampening=dampening,
                        weight_decay=weight_decay, nesterov=nesterov,
                        maximize=maximize)
        if nesterov and (momentum &lt;= 0 or dampening != 0):
            raise ValueError(
                "Nesterov momentum requires a momentum and zero dampening")

        super(SGD, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(SGD, self).__setstate__(state)
        for group in self.param_groups:
            group.setdefault('nesterov', False)
            group.setdefault('maximize', False)

    def _init_group(self, group, params_with_grad, d_p_list, momentum_buffer_list):
        has_sparse_grad = False

        for p in group['params']:
            if p.grad is not None:
                params_with_grad.append(p)
                d_p_list.append(p.grad)
                if p.grad.is_sparse:
                    has_sparse_grad = True

                state = self.state[p]
                if 'momentum_buffer' not in state:
                    momentum_buffer_list.append(None)
                else:
                    momentum_buffer_list.append(state['momentum_buffer'])

        return has_sparse_grad

    @torch.no_grad()
    def step(self, closure=None):
        """Performs a single optimization step.

        Args:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            params_with_grad = []
            d_p_list = []
            momentum_buffer_list = []

            has_sparse_grad = self._init_group(group, params_with_grad, d_p_list, momentum_buffer_list)

            # Calling my custom C++ function
            sgd_cpu.my_sgd_function(params_with_grad,
                                    d_p_list,
                                    momentum_buffer_list,
                                    weight_decay=group['weight_decay'],
                                    momentum=group['momentum'],
                                    lr=group['lr'],
                                    dampening=group['dampening'],
                                    nesterov=group['nesterov'],
                                    maximize=group['maximize'],
                                    has_sparse_grad=has_sparse_grad)

            # update momentum_buffers in state
            for p, momentum_buffer in zip(params_with_grad, momentum_buffer_list):
                state = self.state[p]
                state['momentum_buffer'] = momentum_buffer

        return loss
</code></pre>
<p>The code runs without errors. The problem is that setting the parameter <code>momentum</code> in <code>mySGD</code> has no effect at all. As far as I know, the problem is with the update of the <code>momentum_buffer_list</code>. At the beginning of the training, this parameter is an empty list that should be updated with derivatives <code>d_p_list</code>. This is done in the C code when</p>
<pre><code class="lang-auto">momentum_buffer_list[i] = buf;
</code></pre>
<p>However, when turning back to Python, the same parameter seems to be the previous list of <code>None</code>s, so the <code>state['momentum_buffer']</code> is not updated (I could check this by printing the variables both inside the C code, after assignation, and again in Python, after execution of <code>my_sgd_function</code>).</p>
<p>How could I solve this?<br>
This is a naive example, but my intention is to define a new optimizer as a C++ extension that updates internal parameters as in SGD or Adam.</p>
<p>I realized that in <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/api/src/optim/sgd.cpp" rel="noopener nofollow ugc">pytorch/sgd.cpp at master · pytorch/pytorch (github.com)</a>, the <code>state</code> of the optimizer is used for <code>momentum_buffer</code> update, <code>state-&gt;momentum_buffer(buf);</code>, but I don’t know how to apply this in my case, or if I should use that in a C++ extension.</p>
<p>Any suggestion that helps me out here will be thankful.</p>

Raul_Murillo: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> any insight on how to tackle this?</p>

ptrblck: <p>I don’t know why your code is currently failing and why the <code>momentum_buffer_list</code> is not being updated. It might be easier to write your experiments in Python or did you see significant performance differences in using the C++ implementation?<br>
Note that both would go through the dispatching mechanism and would use the same kernel calls. To reduce the number of kernel calls (if this is adding to the overhead), you could check the <code>fused</code> implementations of a few optimizers (such as <code>Adam</code>) and use e.g. <code>torch._foreach_sub_</code>.</p>

Raul_Murillo: <p>I was able to solve this problem by returning and reassigning the <code>momentum_buffer_list</code> as follows:</p>
<p>In C++ code:</p>
<pre><code class="lang-cpp">std::vector&lt;std::optional&lt;torch::Tensor&gt;&gt; sgd_function(...)
{
...
return momentum_buffer_list;
}
</code></pre>
<p>and then in the Python step function:</p>
<pre><code class="lang-python">momentum_buffer_list = sgd_cpu.sgd(...)
</code></pre>
<p>This is strange to me since Python List objects are mutable, so any function with the list as a parameter should modify it. However, maybe it is because of the C++ extension, or the declaration of the list as <code>std::vector&lt;std::optional&lt;torch::Tensor&gt;&gt; &amp;momentum_buffer_list</code>, it is not mutable in this case.<br>
I could not find a lot of information about this, but it seems that the problem was with Python mutable/immutable objects, rather than an incorrect use of libtorch.</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> the main reason for using C++ extension is not due to the overhead, but cause I want to modify the optimizer using certain operations defined in a third-party C++ library, so porting that to Python is not an option.</p>