Pytorch freeze part of the layers

In PyTorch we can freeze the layer by setting the requires_grad to False. The weight freeze is helpful when we want to apply a pretrained model.

Here I’d like to explore this process.

Build a toy model

If we run it on the terminal, we can explore the parameters

Output the parameters

Set requires_grad to False

Set requires_grad back to True

Freeze part of the parameter

For example, only freeze the fc1 layer.

get params keys

set related layer’s require grads to False (a naive way)

A better way

We can identify the parameter by name[2]:

Then we can filter and control the requires_grad by filtering through the parameter names

Last one more step

We haven’t done yet as even the required grad is set to False, we still can update the weights

So we should filter the parameters to only those requires_grad ones by using this code[1]

Quick summary

we can use

  • net.state_dict() to get the key information of all parameters and we can print it out to help us figure out which layers that we want to freeze
  • If we know our target layer to be frozen, we can then freeze the layers by filtering by names

Thanks for reading and I hope it helps.






Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store