Applying optimizer to a slice of a variable

Ryan_Devera: <p>Hey all,</p>
<p>I have a variable <strong>s</strong> that I would like to slice and then update the the variable <strong>s</strong> only with the sliced portion using an optimizer. Here is my example:</p>
<pre><code class="lang-auto">s = torch.ones(4, 1, requires_grad=True)
optimizer = optim.AdamW([s], lr=1)

for i in range(2):
    optimizer.zero_grad()
    print(i)
    sub_s = s[(2 * i):(2 * (i + 1))]
    sub_s.retain_grad()
    
    f = torch.sum(2 * sub_s)
    f.backward()
    optimizer.step()
</code></pre>
<p>So in the first iteration I get these outputs:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; s
tensor(
[[1.],
 [1.],
 [1.],
 [1.]],
 requires_grad=True
)
&gt;&gt;&gt; s.grad
tensor(
[[2.],
 [2.],
 [0.],
 [0.]]
)
</code></pre>
<p>Then I look at <strong>s</strong> after the optimizer step and unfortunately it is not the update I was hoping for. I actually realize this is because I am using <code>AdamW</code> and even though the gradients are zero the momentum updates are being applied.</p>
<pre><code class="lang-auto">&gt;&gt;&gt; s (after optimizer step)
tensor(
[[-0.0100],
 [-0.0100],
 [ 0.9900],
 [ 0.9900]],
requires_grad=True
)
</code></pre>
<p>If I change the code to, where I take the step outside of the loop and I remove the <code>optimizer.zero_grad()</code> I get the correct updates. However, this is not how I’ve seen typical loops written in <code>torch</code> and also this is a simple example I am wanting to train a model + update slices of variables.</p>
<pre><code class="lang-auto">for i in range(2):
    print(i)
    sub_s = s[(2 * i):(2 * (i + 1))]
    sub_s.retain_grad()
    
    f = torch.sum(2 * sub_s)
    f.backward()

optimizer.step() (step outside of the loop)
</code></pre>
<p>Are there any suggestions on how to do this correctly?</p>

ptrblck: <p>There a are a few approaches such as:</p>
<ul>
<li>Creating different “sub-tensors” and using e.g. <code>torch.cat</code> to create the parameter before using it via the functional API as described <a href="https://discuss.pytorch.org/t/want-to-optimize-only-some-variables-but-getting-cant-optimize-a-non-leaf-tensor-error/139714/6">here</a>.</li>
<li>Restoring the “frozen” part of the parameter after each parameter update.</li>
<li>Guaranteeing that the “frozen” part of the parameter was never and will never be updated in which case even optimizers with running stats or momentum should not update it since past updates were all zero.</li>
</ul>

Ryan_Devera: <p>Hey <a class="mention" href="/u/ptrblck">@ptrblck</a> thank you so much for the suggestion! I feel like I am really close here is a snippet based on your recommendation.</p>
<pre><code class="lang-auto"># Create sub-tensors for the problem
sub_tensors = [torch.ones(2, 1, requires_grad=True) for _ in range(2)]
s = torch.cat(sub_tensors)

# Create a list for the optimizer
sub_tensor_optimizer = [{'params': weights, 'lr': 0.1} for weights in sub_tensors]
optimizer = optim.AdamW(sub_tensor_optimizer)

for i in range(2):
    optimizer.zero_grad()

    # Enable weights to be trainable only for this slice
    # Otherwise turn them off
    for index, weights in enumerate(sub_tensors):
        weights.requires_grad = False
        weights.grad = None
        if i == index:
            weights.requires_grad = True
    
    index = range((2 * i), (2 * (i + 1)))
        
    sub_s = s[index]
    sub_s.retain_grad()
    
    f = torch.sum(2 * sub_s)
    f.backward()

    optimizer.step()
    print(sub_tensors)
</code></pre>
<p>Output</p>
<pre><code class="lang-auto">[tensor([[0.8990], [0.8990]], requires_grad=True), tensor([[1.], [1.]])]
[tensor([[0.8311], [0.8311]]), tensor([[0.8990], [0.8990]], requires_grad=True)]
</code></pre>
<p>The first output look right in the first iteration. I had ‘Frozen’ the second slice of weights and there was no update to the variable. However, even after I ‘Froze’ the first variable in the second iteration there was still some update applied. Any guidance would greatly appreciated!</p>
<p><strong>Update</strong> I also added <code>weights.grad = None</code> after <code>weights.requires_grad = False</code> and now it works the way I expected. Not sure why <code>optimizer.zero_grad()</code> does not take care of that?</p>

ptrblck: <aside class="quote no-group" data-username="Ryan_Devera" data-post="3" data-topic="170056">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ryan_devera/48/53463_2.png" class="avatar"> Ryan_Devera:</div>
<blockquote>
<p><strong>Update</strong> I also added <code>weights.grad = None</code> after <code>weights.requires_grad = False</code> and now it works the way I expected. Not sure why <code>optimizer.zero_grad()</code> does not take care of that?</p>
</blockquote>
</aside>
<p>Calling <code>optimizer.zero_grad()</code> will fill the <code>.grad</code> attribute of each provided parameter with <code>torch.zeros_like(param)</code> and will not delete the <code>.grad</code> attribute by setting it to <code>None</code>.<br>
Use <code>optimizer.zero_grad(set_to_none=True)</code> to delete the <code>.grad</code> attribute instead.</p>