Will freezing an intermediate block influence the gradient flow?

username: ProkaryMonster
cooked: <p>The problem I’m facing is that I want to insert a small pre-trained model to an existing model to do something like features enhancement.<br>
Whereas I want to know if the freezing operation (setting the requires_grad flag of parameters to False) will influence the gradient calculation especially for the layers before the inserted block.</p>
<pre><code class="lang-auto">def __init__(self, block, num_blocks, num_classes=10):
    ...
    ce = MGN_FIN()
   
    # load CE parameters
    ce = torch.nn.DataParallel(ce)
    ckp = torch.load('...')
    ce.load_state_dict(ckp['model'])

    # freeze CE layer's weights
    for name, p in ce.module.named_parameters():
        p.requires_grad = False

    # construct CE module
    self.CE = nn.Sequential(
        ce.module.block1,
        ce.module.block2,
        ce.module.block3,
    )  

...
def forward(self, x):
    out = F.relu(self.bn1(self.conv1(x)))
    out = self.CE(out)

    out = self.layer1(out)
    out = self.CE(out)

    out = self.layer2(out)
    out = self.layer3(out)
    out = F.avg_pool2d(out, out.size()[3])
    out = out.view(out.size(0), -1)
    out = self.linear(out)
    return out

</code></pre>
<p>For example:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/f/1/f1851466d39c65a927258e086196a508f068a9ec.png" alt="image" data-base62-sha1="ysAbVeiJfB5n9sID2SWIKbkB2YA" width="344" height="225"><br>
Any help is appreciated!</p>

username: iffiX
cooked: <p>A cleaner solution would be:</p>
<pre><code class="lang-auto">out = F.relu(self.bn1(self.conv1(x)))
with torch.no_grad():
    out = self.CE(out)
</code></pre>
<p>However, I see that you are using the CE layer several times, and this solution is not feasible since gradient will not flow through your second CE layer, and you will not be able to train <code>self.layer1</code>, therefore, I suggest you either:</p>
<ol>
<li>
<p>manually clear the <code>.grad</code> property and set them to <code>None</code> by iterating on <code>self.CE.parameters()</code> after <code>loss.backward()</code> and before calling <code>optimizer.step()</code></p>
</li>
<li>
<p>restrict the range of parameters passed to the optimizer, pesudo code is:</p>
</li>
</ol>
<pre><code class="lang-auto">optimizer([p for p in your_model.parameters() if p not in your_model.ce.parameters()],
                lr=learning_rate)
</code></pre>
<p>These two ways will not create problems stated in your question.</p>
<p>I am not sure whether setting “requires_grad” to <code>False</code> will interfere with the gradient calculation, may be do a simple experiment with two linear layers yourself?</p>

username: googlebot
cooked: <p>About chain rule, gradients of input transformations wrt input should not be ignored with freezing, i.e. any module can be viewed as some chainable transformation f(x,p) (p - local parameters) and freezing implies ignoring p gradients: f(x,p.detach())</p>

username: ProkaryMonster
cooked: <p>Thanks for your help!</p>
<p>I’ve done the following simple experiment:</p>
<pre><code class="lang-auto">
class Demo(nn.Module):

    def __init__(self):
        super(Demo, self).__init__()
        self.layer1 = nn.Linear(1, 1, bias=False)
        self.layer2 = nn.Linear(1, 1, bias=False)
        self.layer3 = nn.Linear(1, 1, bias=False)

        self.layer1.weight = nn.Parameter(torch.ones([1, 1]).float())
        self.layer2.weight = nn.Parameter(torch.ones([1, 1]).float() * 2)
        self.layer3.weight = nn.Parameter(torch.ones([1, 1]).float() * 3)

        # self.layer1.register_forward_hook(lambda _, x_in, x_out: print('layer1', x_in, x_out))
        # self.layer2.register_forward_hook(lambda _, x_in, x_out: print('layer2', x_in, x_out))
        # self.layer3.register_forward_hook(lambda _, x_in, x_out: print('layer3', x_in, x_out))

        self.layer1.register_backward_hook(lambda _, grad_in, grad_out: print('layer1', grad_in, grad_out))
        self.layer2.register_backward_hook(lambda _, grad_in, grad_out: print('layer2', grad_in, grad_out))
        self.layer3.register_backward_hook(lambda _, grad_in, grad_out: print('layer3', grad_in, grad_out))

        # self.layer2.weight.requires_grad = False

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)

        return out


if __name__ == '__main__':
    model = Demo()
    optimizer = optim.SGD(params=model.parameters(), lr=0.1)
    optimizer.zero_grad()

    print(model.layer1.weight, model.layer2.weight, model.layer3.weight)
    # for i in range(10):
    x = torch.ones((1, 1))
    out = model(x)
    out.backward()
    optimizer.step()
    optimizer.zero_grad()
    print(model.layer1.weight, model.layer2.weight, model.layer3.weight)
</code></pre>
<p>The result is:</p>
<pre><code class="lang-auto"># weights of linear layers
tensor([[1.]], requires_grad=True) 
tensor([[2.]], requires_grad=True) 
tensor([[3.]], requires_grad=True)

# for a linear layer: y = Wx
# gradients of backward_hook  ((gradient for W, gradient for x), gradient for y)
layer3 (tensor([[3.]]), tensor([[2.]])) (tensor([[1.]]),)
layer2 (tensor([[6.]]), tensor([[3.]])) (tensor([[3.]]),)
layer1 (None, tensor([[6.]])) (tensor([[6.]]),)

# updated weights of linear layers 
tensor([[0.4000]], requires_grad=True)
tensor([[1.7000]], requires_grad=True) 
tensor([[2.8000]], requires_grad=True)
</code></pre>
<p>If we set   <code>self.layer2.weight.requires_grad = False</code>, we’ll  get:</p>
<pre><code class="lang-auto">tensor([[1.]], requires_grad=True)
tensor([[2.]]) 
tensor([[3.]], requires_grad=True)

layer3 (tensor([[3.]]), tensor([[2.]])) (tensor([[1.]]),)
layer2 (tensor([[6.]]), None) (tensor([[3.]]),)
layer1 (None, tensor([[6.]])) (tensor([[6.]]),)

tensor([[0.4000]], requires_grad=True) 
tensor([[2.]])
tensor([[2.8000]], requires_grad=True)

</code></pre>
<p>We can see from this that setting the <code>requires_grad </code> flag to <code>False</code> will only set the gradient with respect to the weight of the layer to <code>None</code>,  but the gradient wrt the <code>x</code> will be calculated and preserved for the shallower layer, which can be used for gradient descending correctly.</p>

username: iffiX
cooked: <p>Good experiment. Now it proves that <code>requires_grad=False</code> will calculate gradients with respect to that parameter (layer2.weight) but will not set gradient on it, so that gradient can still flow to layer1, but layer2 is not affected.</p>
<p>There is one tiny error in your comment:</p>
<pre><code class="lang-diff">-# gradients of backward_hook  ((gradient for W, gradient for x), gradient for y)
+# gradients of backward_hook  ((gradient for x, gradient for W), gradient for y)
</code></pre>
<p>And <a class="mention" href="/u/googlebot">@googlebot</a> <code>detach()</code> is definetely inappropriate here, since it will detach the whole remaining chain from backward computation completely, so layer1 will not receive any gradients:</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch.optim as optim
import torch

class Demo(nn.Module):

    def __init__(self):
        super(Demo, self).__init__()
        self.layer1 = nn.Linear(1, 1, bias=False)
        self.layer2 = nn.Linear(1, 1, bias=False)
        self.layer3 = nn.Linear(1, 1, bias=False)

        self.layer1.weight = nn.Parameter(torch.ones([1, 1]).float())
        self.layer2.weight = nn.Parameter(torch.ones([1, 1]).float() * 2)
        self.layer3.weight = nn.Parameter(torch.ones([1, 1]).float() * 3)

        # self.layer1.register_forward_hook(lambda _, x_in, x_out: print('layer1', x_in, x_out))
        # self.layer2.register_forward_hook(lambda _, x_in, x_out: print('layer2', x_in, x_out))
        # self.layer3.register_forward_hook(lambda _, x_in, x_out: print('layer3', x_in, x_out))

        self.layer1.register_backward_hook(lambda _, grad_in, grad_out: print('layer1', grad_in, grad_out))
        self.layer2.register_backward_hook(lambda _, grad_in, grad_out: print('layer2', grad_in, grad_out))
        self.layer3.register_backward_hook(lambda _, grad_in, grad_out: print('layer3', grad_in, grad_out))

        # self.layer2.weight.requires_grad = False

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out).detach()
        out = self.layer3(out)

        return out


if __name__ == '__main__':
    model = Demo()
    optimizer = optim.SGD(params=model.parameters(), lr=0.1)
    optimizer.zero_grad()

    print(model.layer1.weight, model.layer2.weight, model.layer3.weight)
    # for i in range(10):
    x = torch.ones((1, 1))
    out = model(x)
    out.backward()
    optimizer.step()
    optimizer.zero_grad()
    print(model.layer1.weight, model.layer2.weight, model.layer3.weight)
</code></pre>
<p>will print:</p>
<pre><code class="lang-auto">Parameter containing:
tensor([[1.]], requires_grad=True) 
Parameter containing:
tensor([[2.]], requires_grad=True) 
Parameter containing:
tensor([[3.]], requires_grad=True)
layer3 (None, tensor([[2.]])) (tensor([[1.]]),)
Parameter containing:
tensor([[1.]], requires_grad=True) 
Parameter containing:
tensor([[2.]], requires_grad=True) 
Parameter containing:
tensor([[2.8000]], requires_grad=True)
</code></pre>
<p>gradient of layer1 is <strong>not calculated at all</strong>. therefore, <code>requires_grad=True</code> is a correct solution.</p>

username: googlebot
cooked: <aside class="quote no-group" data-username="iffiX" data-post="5" data-topic="88859">
<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/iffix/48/24443_2.png" class="avatar"> iffiX:</div>
<blockquote>
<p>And <a class="mention" href="/u/googlebot">@googlebot</a> <code>detach()</code> is definetely inappropriate here,</p>
</blockquote>
</aside>
<p>Yeah, I just used that as notation for functional style equivalent. In OOP style, weight is not a function argument, so for nn.Parameters you usually disable gradient with requires_grad.</p>