How to customize the "backward" function to not calculate gradients and not update a selected sub-tensor of a parameter tensor of a layer

eric4: <p>Let’s take the model below :</p>
<pre><code class="lang-auto">import torch.nn as nn
from torch import optim
import torch.nn.functional as F
import torch

class NeuralNetwork(nn.Module):

    def __init__(self,input_size,output_size):
        super(NeuralNetwork, self).__init__()
        self.linear1=nn.Linear(input_size, 10)
        self.linear2=nn.Linear(10, output_size)

    def forward(self, x):

        x = self.linear1(x)
        x = nn.functional.relu(x)

        x = self.linear2(x)

        return x

</code></pre>
<pre><code class="lang-auto">model=NeuralNetwork(5,2)
model.state_dict()["linear1.weight"].shape

# output :
# torch.Size([10, 5])
</code></pre>
<p>Run random data through the model and optimize the parameters(with just one forward and backward propagation):</p>
<pre><code class="lang-auto">input=torch.randn(200,5)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=1, momentum=0.9)
output=model(input)
loss = criterion(output,torch.randn(200,2))

# update
loss.backward()
optimizer.step()
</code></pre>
<p>for example, for the first layer,I want to update only a sub-tensor of the weights tensor  :</p>
<p>for a given <code>idx</code>, <code>model.state_dict()["linear1.weight"][0:idx]</code> should not be updated,(and apply the same reasoning for other layers,with different <code>idx</code>)</p>
<p>I don’t want to replace the gradients after being calculated by 0 because my goal is to avoid computing them for optimization reasons,</p>
<p>any help would be appreciated</p>

ptrblck: <aside class="quote no-group" data-username="eric4" data-post="1" data-topic="169103">
<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/e/bbe5ce/48.png" class="avatar"> eric4:</div>
<blockquote>
<p>I don’t want to replace the gradients after being calculated by 0 because my goal is to avoid computing them for optimization reasons,</p>
</blockquote>
</aside>
<p>In that case you could create a “fixed” and a “trainable” tensor in a custom linear layer and concatenate them in each <code>forward</code> pass. This would make sure that only the trainable part gets valid gradient and the parameter updates.<br>
<a href="https://discuss.pytorch.org/t/custumize-the-backpropagation-phase-of-a-neural-network-to-ignore-the-update-of-some-parameters/167557/4">This post</a> gives you an example of such a layer and replaces it via <code>torch.fx</code> in another model.</p>

eric4: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>I took a look at this post and from the results presented there, operations like <code>torch.cat</code> can consistently slow down the forward and backward propagation,</p>
<p>Is it the only way to do it,because if I understood correctly, we can’t do withou <code>torch.cat</code> calls</p>
<p>Thanks in advance</p>

ptrblck: <p>No, it’s not the only way as you could also:</p>
<ul>
<li>zero out the gradients assuming your optimizer does not use internal running states to update parameters or</li>
<li>restore the parameters after a full update.</li>
</ul>
<p>However, since you don’t want to compute the gradients for the frozen part of your parameters at all, the I think the <code>torch.cat</code> is the only one meeting your requirements.</p>

eric4: <p>thank you for the clarification <a class="mention" href="/u/ptrblck">@ptrblck</a> ,</p>
<p>do you think it is feasible to improve the code of the post so that training the model with fixed tensors will be faster than training the original model withoud fixed ones(since I want to use it  for optimization reasons)?  it seems to me that the code is already optimal…<br>
I know my question is a bit general, but I want an expert opinion to know if it’s worth digging deeper on this,</p>
<p>Thanks</p>

ptrblck: <p>I think the fastest approach would depend on your actual use case and I would suggest to profile the discussed methods.<br>
In particular, using the <code>torch.cat</code> approach with the <code>out</code> argument could be beneficial if only a small part of the actual gradients should be calculated. On the other hand, you might not see huge benefits of avoiding the computation of the “frozen” gradients if the actual workload is tiny. Since you are concerned about the optimal performance, a profile would be the right approach to see which method would add the most overhead.</p>

eric4: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>  I have a question :</p>
<p>Why is it necessary to use torch.fx <a href="https://discuss.pytorch.org/t/custumize-the-backpropagation-phase-of-a-neural-network-to-ignore-the-update-of-some-parameters/167557/4">here</a> ?</p>
<p>Is it possible to simply replace the selected layer with the customized one which will freeze the desired parameters ( using <code>setattr(model, module_name, new_module)</code>)  without losing the weights in the transition?</p>

ptrblck: <p>You don’t need to use <code>torch.fx</code>, but it might be a convenient way to replace layers “automatically”.<br>
Yes, you could also replace layers manually by replacing the internal attributes.</p>