How the pytorch freeze network in some layers, only the rest of the training?

does this work even if the network has already been trained? Say I have loaded a pre-trained net on X and I want to freeze layer Y (say 2nd layer to make the example concrete). How do I do that exactly?

This snippet may clarify how to do it.

Set requires_grad to false you want to freeze:

# we want to freeze the fc2 layer
net.fc2.weight.requires_grad = False
net.fc2.bias.requires_grad = False

Then set the optimizer like the following:

optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()), lr=0.1)

Alternatively, you can only add the parameters you want to train to the optimizer:
https://discuss.pytorch.org/t/to-require-grad-or-to-not-require-grad/5726

But I think the above method is more straight-forward.

14 Likes

@L0SG What if we want to unfreeze those layers later in the code? For example, I want to fine tune heads first, then tune the other layers as well. We will need to instantiate a new optimizer, right? If so, does it affect the optimization?

2 Likes

optimizer.add_param_group would be what you want. This will unfreeze the previously frozen layer, by adding the parameters to new dict elements of param_groups (list) of the optimizer:

# let's unfreeze the fc2 layer this time for extra tuning
net.fc2.weight.requires_grad = True
net.fc2.bias.requires_grad = True

# add the unfrozen fc2 weight to the current optimizer
optimizer.add_param_group({'params': net.fc2.parameters()})

10 Likes

if you only do:

# we want to freeze the fc2 layer
net.fc2.weight.requires_grad = False
net.fc2.bias.requires_grad = False

without the second part does it still work? i.e. without:

optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()), lr=0.1)
3 Likes

Yes, it does work when you add the parameters with requires_grad=True to the optimizer then setting to False after. You can also find out yourself by commenting out

optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()), lr=0.1)

In the snippet above, since the previous optimizer contains all parameters including the fc2 with the changed requires_grad flag.

Note that the above snippet assumed a common “train => save => load => freeze parts” scenario.

Snippet for recursively freezing a portion of your graph.

def dfs_freeze(model):
    for name, child in model.named_children():
        for param in child.parameters():
            param.requires_grad = False
        dfs_freeze(child)

Then:

optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=opt.lr, amsgrad=True)
2 Likes

@L0SG, When I unfreeze the previously frozen layers and also want to change the learning rate and the Momentum, but the parameters of these layers don’t have ‘momentum_bufer’ in optimizer’s state, do you have any suggestions to this issue?

class YourModel(nn.Module):
    
    def __init__(self, in_features, out_features):
        super(YourModel, self).__init__()
        self.fc = nn.Linear(in_features, out_features)

    def forward(self, x):
        return self.fc(x)

model = YourModel(in_features, out_features)
model.fc.train(False)

Hello I am still confuse about freeze the weight in Pytorch looks like very hard to do. Suppose I want to make a loss function which filtering the loss using the initialized kernel. I am using nn.conv2D to do my job but I don’t want the weight being updated(freeze). The loss function basically the simple network let said the A network is the main network that will be updated, and B Network which is the network for computing the loss function. For this task I enabled the grad for A and disable for B. During the training my program will taking that loss from B, then backpropagate into the main network A (where the weight should be update). However I always ended with this.
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
Anyone have the better way to do this?, can I just use modelB.eval() ? or modelB.Train(False)?

1 Like

You can try setting the learning rate for those to 0. Optimizer support different learning rate for each param groups, it can be adjusted online, look that up.
Note that layers (such that BatchNorm) do the learning in forward phase, so you have to set them to eval() too

2 Likes

If we set the requires_grad to false for a particular layer, do we have to leave it out of the optimizer?
Such as, this ->

optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()), lr=0.1)

Is the benefit in the speedup? Or is it wrong if one does not filter it out of the optimizer?
Thanks

The error from adding already frozen parameters to the optimizer is an explicit design choice to avoid silly mistakes (one could mistakenly add the frozen one and assume that the parameter is being trained). But manually freezing parameters after declaring the optimizer works as intended without re-defining the optimizer. (my own replay above in this thread)

Thanks for the clarification!
As a follow-up to my point regarding the speed up - I am not observing a speedup when I freeze the initial 80% of the network. I expected the training to be faster, since it only has to update 20% and to be lighter, since it only has to store the information to execute a backward for 20% of the network.
Is there a speedup expected in this scenario?

I’m afraid I have no definitive answer for this since I don’t know your exact model setup, but several suggestions:

  1. Every single tensor before the frozen part in the computational graph must also be requires_grad=False so that the frozen subgraph gets excluded in the autograd engine. If there exists any tensor that requires grad, It’ll need all the backward pass of the graph anyway.
  2. I would check which part of your model is a major bottleneck in speed. For example, if the unfrozen parts may contain parameters that require heavy-lifting like this example. Or, the bottleneck would be unrelated to the model itself, such as loading your data that overshadows the model speed.
  3. Make sure that you use cuda.synchronize() for properly measuring the speed.
1 Like

Have you solved this problem?

@lugiavn’s method should work.
As long as you need to compute d(B)/d(params of A), you have to backpropagate gradient along paths in B to A. requires_grad attribute has to be set for B. Maybe writing custom backward function for B will be more a efficient way, it’ll be constant as there is no parameter updating.

I just experimenting with several things and I found that the B network weight should be requires_grad = False, but the tensor should be always requires_grad = True. and for the optimizer, just optimize the A networks. It works fine for me. if you disable the grad, in the end, you need to enable again before calculating the loss, because the backpropagation need the gradient history. I don’t know the exact solution but it works fine for me and exactly what I want.

Thanks for your detailed explanation.

But this does not make sense to me.
Could you have a further explanation on network weight’s requires_grad and tensor’s requires_grad?

Suppoose we have a network B part we want to freeze.

for param in networkB.conv1.parameters():
    param.requires_grad = False

For the tensor, we can set it while creating the tensor. you can see details in here.
x = torch.tensor([1], requires_grad=True)

1 Like