Custom optimizer replacing param.add_ in SGD

lnair: <p>Hello,</p>
<p>I’m trying to implement a custom SGD optimizer, by changing the way in which <code>d_p * alpha</code> gets added to <code>param</code>. I started out with the following custom implementation which copies SGD code from: <a href="https://pytorch.org/docs/stable/_modules/torch/optim/sgd.html#SGD" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/_modules/torch/optim/sgd.html#SGD</a> with a small change in <code>_single_tensor_sgd()</code>:</p>
<pre><code class="lang-auto">def _single_tensor_sgd(
    params,
    d_p_list,
    momentum_buffer_list,
    *,
    weight_decay,
    momentum,
    lr,
    dampening,
    nesterov,
    maximize,
    has_sparse_grad
):
    for i, param in enumerate(params):

        d_p = d_p_list[i]
        if weight_decay != 0:
            d_p = d_p.add(param, alpha=weight_decay)

        if momentum != 0:
            buf = momentum_buffer_list[i]

            if buf is None:
                buf = torch.clone(d_p).detach()
                momentum_buffer_list[i] = buf
            else:
                buf.mul_(momentum).add_(d_p, alpha=1 - dampening)

            if nesterov:
                d_p = d_p.add(buf, alpha=momentum)
            else:
                d_p = buf

        alpha = lr if maximize else -lr

        # Changed this part
        #param.add_(d_p, alpha=alpha)  # &lt;-- Matches pytorch
        param = param + d_p * alpha   # &lt;-- Doesn't match pytorch
</code></pre>
<p>To check whether this change makes any difference, I tested it against Pytorch’s implementation as shown below:</p>
<pre><code class="lang-auto">import torch
# MySGD is the same as torch SGD except the _single_tensor_sgd function
from test_SGD import MySGD
import copy

class SimpleLinear(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = torch.nn.Linear(in_features, out_features)

    def forward(self, x):
        return self.linear(x)

in_features = 10
out_features = 50

# Create two equivalent models
torch_model = SimpleLinear(in_features, out_features)
my_model = copy.deepcopy(torch_model)
my_model.load_state_dict(torch_model.state_dict())

torch_model.train()
my_model.train()

lr = 0.001

# My optimizer
my_optim = MySGD(my_model.parameters(), lr)
# Pytorch optimizer
torch_optim = torch.optim.SGD(torch_model.parameters(), lr)

# Mock inputs
inputs = torch.randn((2, in_features))

# Performs "training" step
def get_my_opt_state():
    torch.manual_seed(10)
    y = my_model(inputs)
    y.sum().backward()
    my_optim.step()
    my_model.zero_grad()

def get_torch_opt_state():
    torch.manual_seed(10)
    y = torch_model(inputs)
    y.sum().backward()
    torch_optim.step()
    torch_model.zero_grad()

# Compute error between torch and my weights
def compute_abs_error(x, y):
    return torch.sum(torch.abs(x - y))

# Do one mock epoch
for epoch in range(1):
    get_torch_opt_state()
    get_my_opt_state()

# Compare resultant weights
torch_w = torch_model.linear.weight
my_w = my_model.linear.weight

error = compute_abs_error(torch_w, my_w)
print("Error: ", error)
</code></pre>
<p>However, with <code>param.add_(d_p, alpha=alpha)</code> there is zero error. Once changed to <code>param = param + d_p * alpha</code> it gives me a non-zero error. What could be causing this difference? Or, is my testing incorrect? I would like to write a custom way of doing the update step, so I want to make sure I don’t cause the same issue there.</p>

KFrank: <p>Hi lnair!</p>
<aside class="quote no-group" data-username="lnair" data-post="1" data-topic="157297" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/91b2a8/48.png" class="avatar"> lnair:</div>
<blockquote>
<pre><code class="lang-auto">        # Changed this part
        #param.add_(d_p, alpha=alpha)  # &lt;-- Matches pytorch
        param = param + d_p * alpha   # &lt;-- Doesn't match pytorch
</code></pre>
</blockquote>
</aside>
<p><code>param = param + d_p * alpha</code> creates a new tensor and sets the<br>
python reference <code>param</code> to refer to the new tensor, rather than modifying<br>
the tensor you wish to train.  Conversely, <code>param.add_(d_p, alpha=alpha)</code><br>
modifies the desired tensor (“adds in place”).</p>
<p>One approach that would keep much of your code unchanged would be:</p>
<pre data-code-wrap="python"><code class="lang-python">with torch.no_grad():
    param.copy_ (param + d_p * alpha)
</code></pre>
<p>To illustrate the issue of setting <code>param</code> to refer to a new tensor, consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.10.2'
&gt;&gt;&gt; param = torch.ones (2, 3)
&gt;&gt;&gt; orig_param = param
&gt;&gt;&gt; _ = param.add_ (torch.ones (2, 3))
&gt;&gt;&gt; param
tensor([[2., 2., 2.],
        [2., 2., 2.]])
&gt;&gt;&gt; orig_param
tensor([[2., 2., 2.],
        [2., 2., 2.]])
&gt;&gt;&gt; param = param + torch.ones (2, 3)
&gt;&gt;&gt; param
tensor([[3., 3., 3.],
        [3., 3., 3.]])
&gt;&gt;&gt; orig_param
tensor([[2., 2., 2.],
        [2., 2., 2.]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

lnair: <p>Thank you for the response! That clears it up for me. There is a still a very small difference in results (order of e-9), but I can work around that.</p>