How can I only train the classifier and freeze rest of the parameters in Pytorch?

Question

I have taken the pretrained model of MoviNet, I have changed the last layer.

This is last parameters of pretrained model that I have taken;

classifier.0.conv_1.conv2d.weight  :  torch.Size([2048, 640, 1, 1])
classifier.0.conv_1.conv2d.bias  :  torch.Size([2048])
classifier.3.conv_1.conv2d.weight  :  torch.Size([600, 2048, 1, 1])
classifier.3.conv_1.conv2d.bias  :  torch.Size([600])

The following are the parameters that I have changed at the last layer;

clfr.0.multi_head.0.head2.0.conv_1.conv2d.weight  :  torch.Size([2048, 640, 1, 1])
clfr.0.multi_head.0.head2.0.conv_1.conv2d.bias  :  torch.Size([2048])
clfr.0.multi_head.0.head1.weight  :  torch.Size([600, 2048, 1, 1])
clfr.0.multi_head.0.head1.bias  :  torch.Size([600])

I want to train only classifier (clfr) based on previous layer weights, and freeze all previous laers in pytorch, can anyone one tell me how canI do this?

Answer 1

You can set layer.requires_grad=False for each layer that you do not wish to train. If it is easier, you can set it to False for all layers by looping through the entire model and setting it to True for the specific layers you have in mind. This is to ensure you have all other layers set to False without having to explicitly figure out which layers those are.

Answer 2

When creating your optimizer, only pass the parameters that you want to update during training. In your example, it could look something like:

optimizer = torch.optim.Adam(clfr.parameters())