Revert optimizer.step()?

Neofytos: <p>Hi everyone,</p>
<p>I am trying to perform the following experiment and I’d like your advice on what’s the best way to implement it in Pytorch. Given a mini-batch, weight gradients dW^{(t)} are computed based on minimizing a loss function. When we step towards that direction, we get our new weights and can calculate the new loss.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/3/03e9272d131d2a3f8f4769fbe31768e980ed73ac.png" alt="image" data-base62-sha1="yAXy4HUuUEmXzVA97yMFJ9X2lS" width="350" height="180"></p>
<p>What I’d like to do is based on the weight gradients (in particular, the vector direction) to step at multiple distances from W^(t), each time computing the new loss, i.e. move across the dotted line as if I am experimenting with different learning rates/step sizes and <em><strong>each time calculate the loss as if we were to step that much</strong></em>:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/e/3eecaf7cb94ce4a906734df13ea69bfd15ca035d.png" data-download-href="https://discuss.pytorch.org/uploads/default/3eecaf7cb94ce4a906734df13ea69bfd15ca035d" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/e/3eecaf7cb94ce4a906734df13ea69bfd15ca035d_2_517x175.png" alt="image" data-base62-sha1="8YEHmTUQdIzKOdGQZ04kwbdmG85" width="517" height="175" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/e/3eecaf7cb94ce4a906734df13ea69bfd15ca035d_2_517x175.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/e/3eecaf7cb94ce4a906734df13ea69bfd15ca035d_2_775x262.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/e/3eecaf7cb94ce4a906734df13ea69bfd15ca035d.png 2x" data-dominant-color="E7E8E5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">779×265 51.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>In other words, to call optimizer.step() but then be able to revert bach each time to the initial state, so that I can then call optimizer.step() but with a bigger step size, and to repeat this multiple times. What would the best way be to implement this?</p>
<p>Diagrams taken from the work of Santurkar et al. (<a href="https://papers.nips.cc/paper/7515-how-does-batch-normalization-help-optimization.pdf" rel="noopener nofollow ugc">https://papers.nips.cc/paper/7515-how-does-batch-normalization-help-optimization.pdf</a>).</p>

ptrblck: <p>One possible approach would be to create a copy of the <code>state_dict</code> before calling the initial <code>step()</code> and then restoring the model to the initial state.<br>
However, this should remove all <code>grad</code> attributes, so you could either store them separately or rerun the loss calculation.</p>

Neofytos: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thanks for replying! Right now, I am doing it a bit differently (I will try your method right after). My initial learning rate was 1, so I’ve set the learning rate to 0.5, and called optimizer.step() 8 times, each time calculating the loss. However, when I try to step “back” by setting the learning rate to -3 and calling optimizer.step(), I don’t get the same loss that I had calculated in the previous step. Do you know why that would be the case?</p>
<p>Step 0, Accuracy 0.7734375, Loss 1.1175949573516846<br>
Step 1, Accuracy 0.8515625, Loss 0.6774735450744629<br>
Step 2, Accuracy 0.85546875, Loss 0.5671432614326477<br>
Step 3, Accuracy 0.841796875, Loss 0.5469013452529907<br>
Step 4, Accuracy 0.83203125, Loss 0.559360921382904<br>
Step 5, Accuracy 0.80078125, Loss 0.5869375467300415<br>
Step 6, Accuracy 0.783203125, Loss 0.6219274401664734<br>
Step 7, Accuracy 0.76953125, Loss 0.6623443365097046</p>
<p>After stepping back with -3 lr: Loss: 0.7005575299263 (it should have been 0.6774735450744629)</p>
<p>I assumed that since I did not call loss.backwards(), the same gradients would have been applied each time. Is that incorrect?</p>

Neofytos: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> When optimizer.step() is called more than once, the GPU usage increases during training until a CUDA out of memory error is thrown. Do you know why this is the case :)? Thank you!</p>

ptrblck: <p>Your approach might work, if you are using a simple optimizer, e.g. <code>SGD</code> without momentum.<br>
When you are using momentum or an optimizer with running estimates, the reversed step might not work out of the box and you might need to look into the applied formula to check, how to revert it.</p>
<p>Here is a small example using <code>SGD</code>:</p>
<pre><code class="lang-python">torch.manual_seed(2809)

# Set model to eval to prevent batchnorm and dropout layers of changing the output
model = models.resnet50().eval()
x = torch.randn(2, 3, 224, 224)
target = torch.randint(0, 1000, (2,))

optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

# Dummy update steps
out = model(x)
loss = criterion(out, target)
loss.backward()
print('Initial loss ', loss.item())
optimizer.step()
fc0 = model.fc.weight.detach().clone()

# Get updated loss
out = model(x)
loss = criterion(out, target)
print('Updated loss ', loss.item())
fc1 = model.fc.weight.detach().clone()

# Use negative lr
optimizer.param_groups[0]['lr'] = -1. * optimizer.param_groups[0]['lr']
optimizer.step()
out = model(x)
loss = criterion(out, target)
print('Reverted loss ', loss.item())
fc2 = model.fc.weight.detach().clone()
</code></pre>
<p><code>optimizer.step()</code> should not increase the memory usage. Are you sure this line causes the OOM issue?<br>
Usually you would run out of memory, if you are storing some tensors, which are not detached from the computation graph.</p>

Neofytos: <p>If I don’t call loss.backwards after optimizer.step(), the GPU usage keeps on increasing.</p>
<pre><code class="lang-auto">for step_index in range(8):
    optimizer.step()
    outputs = model(X)
    loss = loss_function(outputs, y)
    loss.backward() // If I comment out, GPU usage increases until OOM error

    if step_index == 1:
        to_load_state = {
                'model': model.state_dict()
        }

</code></pre>

G.M: <p>Did u forget to call <code>optimizer.zero_grad()</code>? I think that if u don’t call it the gradients will just stack up and increase the memory usage.</p>

ptrblck: <p>Are you storing the loss or gradients somehow?<br>
The gradients should be accumulated, so the memory footprint should not increase.</p>

Neofytos: <p>You were right. I was storing the gradients. Thank you! It works now <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Note that I used the negative learning rate approach.</p>
<p>Now, the next step is to compute new gradients at each point along the gradient vector line. However, loss.backwards is not an option as that would update the model. Is there any way of computing the current gradients using a function that returns them, rather than changing them in-place?</p>

Yaroslav_Bulatov: <p>You’ll need need to hack it by updating parameters for each evaluation point and then restoring to original values. It’s kind of inherent to structure of torch.nn – since parameters are a state you need to modify state. Parameter=state structure makes it easier for a common cases like SGD, but awkward when you need to try several parameter values per step (ie, line-search or bayesian optimization)</p>

Neofytos: <p>What if I clone the model and do loss.backwards on that model?</p>
<pre><code class="lang-auto">for step_index in range(8):
    optimizer.step()
    temp_model = copy.deepcopy(model)
    temp_outputs = temp_model(X)
    cur_loss = loss_function(temp_outputs, y)
    temp_loss = cur_loss.item()
    cur_loss.backward()
</code></pre>
<p>Would that be an overkill? Also, since the goal is to calculate the L2 norm between these new gradients and the initial gradients, I could then just do:</p>
<pre><code class="lang-auto">    temp_totall2norm = 0
    for i, p1 in enumerate(model.parameters()):
        for j, p2 in enumerate(temp_model.parameters()):
            if i != j:
                continue
        temp_totall2norm += LA.norm(p1.grad.data.cpu().detach().clone() - p2.grad.data.cpu().detach().clone())
</code></pre>
<p>However, when I do run the above, strangely, the LA.norm line seems to be modifying the initial model (I get different results than when I don’t compute the L2 Norm ). Any idea why this is the case?</p>