How to use a layer with gradient but without weight adjustment?

Spen: <p>Is it possible to mark part of the forward pass to only backpropagate the gradient but not to adjust weights?</p>
<p>In the following example code I have a <code>Module</code> that uses only one layer (one set of parameters) but it is used twice in the forward step. During the optimization I would expect the weights to be adjusted twice as well. If I want to only adjust the weights for one of the layer usages, what can I do?</p>
<pre><code class="lang-python">import torch
    
class ExampleModel(torch.nn.Module):
    
    def __init__(self, dim) -&gt; None:
        super(ExampleModel, self).__init__()
        self.linear = torch.nn.Linear(dim, dim)
    
    def forward(self, x):
        out1 = self.linear(x)  # backprop gradients and adjust weights here
        out2 = self.linear(out1)  # only backprop gradients here
        return out2
    
    
# Random input output data for this example
N, D = 64, 100
x = torch.randn(N, D)
y = torch.randn(N, D)
    
model = ExampleModel(D)
    
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.Adam(model.parameters())
    
y_pred = model(x)
loss = criterion(y_pred, y)
    
optimizer.zero_grad()
loss.backward()
optimizer.step()
</code></pre>
<p>The following will not work since with <code>torch.no_grad()</code> no gradient at all is backpropagated:</p>
<pre><code class="lang-python">def forward(self, x):
    out1 = self.linear(x)  # backprop gradients and adjust weights here
    with torch.no_grad():
        out2 = self.linear(out1)  # only backprop gradients here
    return out2
</code></pre>
<p>I can not simply exclude the parameters from the optimization since they need to be optimized in the first part (i.e. <code>out1 = self.linear(x)</code>).<br>
For the same reason I can also not set a learning rate of 0 for these parameters.</p>
<p>What else can I do to achieve this?</p>
<p><em>I asked this question on Stackoverflow but based on other questions I found during my google search it seems that the pytorch community is not very active on Stackoverflow but is active on this forum. Therefore I repost my question here. See my Stackoverflow question here: <a href="https://stackoverflow.com/questions/62487201" rel="nofollow noopener">https://stackoverflow.com/questions/62487201</a></em></p>

ptrblck: <p>The posted answer on StackOverflow seems to be valid.<br>
The second approach using <code>detach()</code> might be easier with the functional API.</p>

Spen: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> That is confusing to me.</p>
<p>What is the difference between <code>torch.no_grad()</code> and to wrap the line with <code>.requires_grad_(False) -&gt; .requires_grad_(True)</code>?<br>
That seems equivalent to each other.</p>
<p>With <code>torch.no_grad()</code> I get a runntime exception since there is no gradient:</p>
<pre data-code-wrap="python"><code class="lang-python">def forward(self, x):
    out1 = self.linear(x)  # backprop gradients and adjust weights here
    with torch.no_grad():
        out2 = self.linear(out1)  # only backprop gradients here
    return out2
</code></pre>
<blockquote>
<p>RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</p>
</blockquote>
<p>By wrapping it with <code>.requires_grad_(False) -&gt; .requires_grad_(True)</code> it seems to run through but I don’t get why. For me that looks identical to the first solution.</p>
<pre data-code-wrap="python"><code class="lang-python">def forward(self, x):
    out1 = self.linear(x)  # backprop gradients and adjust weights here
    self.linear.requires_grad_(False)
    out2 = self.linear(out1)  # only backprop gradients here
    self.linear.requires_grad_(True)
    return out2
</code></pre>
<p>I have the same confusion with <code>.detach</code>. If I would detach it, then how should the gradient be backpropagated?</p>
<p>Maybe the confusion is on the terminology here but I would suspect <code>no_grad</code>, <code>requires_grad(False)</code> and <code>detach</code> to block gradient propagation. Maybe a feature request for a <code>no_weight_adjustment</code> / <code>without_weight_adjustment</code> parameter or environment makes sense?</p>
<p>How should the detach solution work? Could you give an example or add a link to a source where something similar is described?</p>

ptrblck: <p><code>torch.no_grad()</code> is more aggressive, as it Autograd won’t track any operations in the block, so that <code>out2</code> won’t have a valid <code>.grad_fn</code> and you thus won’t be able to call <code>backward</code> on the calculated loss.<br>
<code>.requires_grad</code> will disable the grad calculation for this particular parameter, while the output of the operation will still have a <code>.grad_fn</code>, if one input had previously a <code>.grad_fn</code> (<code>out1</code> in your example).</p>
<p>Here is an example for all three workflows:</p>
<pre><code class="lang-python">class ExampleModel(torch.nn.Module):
    
    def __init__(self, dim) -&gt; None:
        super(ExampleModel, self).__init__()
        self.linear = torch.nn.Linear(dim, dim)
    
    def forward(self, x, mode):
        out1 = self.linear(x)  
        if mode == 'original':
            out2 = self.linear(out1)
        elif mode == 'detach':
            out2 = F.linear(out1, self.linear.weight.detach(), self.linear.bias.detach())
        elif mode == 'req_grad':
            self.linear.requires_grad_(False)
            out2 = self.linear(out1)
            self.linear.requires_grad_(True)
        else:
            raise RuntimeError
        return out2
    
torch.manual_seed(2809)
# Random input output data for this example
N, D = 64, 100
x = torch.randn(N, D)
y = torch.randn(N, D)
    
model = ExampleModel(D)
criterion = torch.nn.MSELoss(reduction='sum')
    
# Original
y_pred = model(x, mode='original')
loss = criterion(y_pred, y)
loss.backward()

weight_grad_ref = model.linear.weight.grad.clone()
bias_grad_ref = model.linear.bias.grad.clone()
model.zero_grad()

# Detach
y_pred = model(x, mode='detach')
loss = criterion(y_pred, y)
loss.backward()

weight_grad_detach = model.linear.weight.grad.clone()
bias_grad_detach = model.linear.bias.grad.clone()
model.zero_grad()

#Req_grad
y_pred = model(x, mode='req_grad')
loss = criterion(y_pred, y)
loss.backward()

weight_grad_req_grad = model.linear.weight.grad.clone()
bias_grad_req_grad = model.linear.bias.grad.clone()
model.zero_grad()

# Compare
print((weight_grad_ref - weight_grad_detach).abs().max())
print((weight_grad_detach - weight_grad_req_grad).abs().max())
</code></pre>
<p>As you can see, the “detach” and “requires_grad” approach yield the same output.</p>

Spen: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thanks for the explanation and the extended example.</p>
<p>Okay, I get how the detach solution was meant here. Of course in my real world problem the model is not just a linear layer. I guess with a detach solution I would need to redefine my whole model with functions. I guess wrapping with <code>requires_grad</code> is easier then. Is there a downside to using <code>requires_grad</code> compared to functions (aka the detach solution)?</p>
<p>I still think there could be some improvement of pytorch here. The difference you mentioned between <code>no_grad</code> and <code>requires_grad</code> is not clear and I think a specific parameter or environment for the “propagate gradient but don’t adjust weights” use-case would be helpful. It would make this much more explicit (I prefer explicitness).<br>
What is the right place to propose such a feature request? Do you see any chance for such a feature request to succeed?</p>

ptrblck: <aside class="quote no-group" data-username="Spen" data-post="5" data-topic="86190">
<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/spen/48/25841_2.png" class="avatar"> Spen:</div>
<blockquote>
<p>What is the right place to propose such a feature request?</p>
</blockquote>
</aside>
<p>You could create a feature request on <a href="https://github.com/pytorch/pytorch/issues">GitHub</a> and explain the shortcomings you are seeing as well as a proposal (if you have a good idea how it should look).</p>
<aside class="quote no-group" data-username="Spen" data-post="5" data-topic="86190">
<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/spen/48/25841_2.png" class="avatar"> Spen:</div>
<blockquote>
<p>Do you see any chance for such a feature request to succeed?</p>
</blockquote>
</aside>
<p>It depends on the request and if it would make all use cases more explicit or could yield ambiguous behavior in specific use cases.</p>

noanabeshima: <p>If you want to repeatedly do this in your code, possibly with multiple module gradients disabled at once, you could use a context manager that temporarily disables the gradients of the parameters in modules you specify.</p>
<p>Here it is added alongside the other examples from ptrblck:</p>
<pre><code class="lang-auto">class disable_module_grad:
    def __init__(self, *modules):
        self.modules = modules
        self.cache = {}
    def __enter__(self):
        for m in self.modules:
            self.cache[m] = {}
            for p in m.parameters():
                self.cache[m][p] = p.requires_grad
                p.requires_grad = False
    
    def __exit__(self, *args):
        for m in self.modules:
            for p in m.parameters():
                p.requires_grad = self.cache[m][p]
        self.cache = {}


class ExampleModel(torch.nn.Module):
    
    def __init__(self, dim) -&gt; None:
        super(ExampleModel, self).__init__()
        self.linear = torch.nn.Linear(dim, dim)
    
    def forward(self, x, mode):
        out1 = self.linear(x)  
        if mode == 'original':
            out2 = self.linear(out1)
        elif mode == 'detach':
            out2 = F.linear(out1, self.linear.weight.detach(), self.linear.bias.detach())
        elif mode == 'req_grad':
            self.linear.requires_grad_(False)
            out2 = self.linear(out1)
            self.linear.requires_grad_(True)
        elif mode == 'context_manager':
            with disable_module_grad(self.linear):
                out2 = self.linear(out1)
        else:
            raise RuntimeError
        return out2
    
torch.manual_seed(2809)
# Random input output data for this example
N, D = 64, 100
x = torch.randn(N, D)
y = torch.randn(N, D)
    
model = ExampleModel(D)
criterion = torch.nn.MSELoss(reduction='sum')
    
# Original
y_pred = model(x, mode='original')
loss = criterion(y_pred, y)
loss.backward()

weight_grad_ref = model.linear.weight.grad.clone()
bias_grad_ref = model.linear.bias.grad.clone()
model.zero_grad()

# Detach
y_pred = model(x, mode='detach')
loss = criterion(y_pred, y)
loss.backward()

weight_grad_detach = model.linear.weight.grad.clone()
bias_grad_detach = model.linear.bias.grad.clone()
model.zero_grad()

# custom
y_pred = model(x, mode='context_manager')
loss = criterion(y_pred, y)
loss.backward()

weight_grad_ctx_manager = model.linear.weight.grad.clone()
bias_grad_ctx_manager = model.linear.bias.grad.clone()
model.zero_grad()

#Req_grad
y_pred = model(x, mode='req_grad')
loss = criterion(y_pred, y)
loss.backward()

weight_grad_req_grad = model.linear.weight.grad.clone()
bias_grad_req_grad = model.linear.bias.grad.clone()
model.zero_grad()

# Compare
print((weight_grad_ref - weight_grad_detach).abs().max())
print((weight_grad_detach - weight_grad_req_grad).abs().max())
print((weight_grad_ctx_manager - weight_grad_detach).abs().max())
</code></pre>