How to freeze part of selected layer(eg nn.Linear()) of a model in Pytorch?

Viewed 31

question: fc = nn.Linear(n,3); I want to freeze the parameters of the third output of the fc when I train this layer.

2 Answers

With monolithic layer you can't. But you can split layer, making a separate Linear for each output channel:

fc1 = nn.Linear(n,1)
fc2 = nn.Linear(n,1)
fc3 = nn.Linear(n,1)

Here you can freeze fc3

I believe it is not possible yet, however the common solution is to set the gradient to zero before backpropagation so the current weights do not change.

With param the parameter of your layer you want to freeze:

param.grad[2,:] = torch.zeros_like(param.grad[2,:])

before calling loss.backward()

Don't forget to set to zero the adequate bias gradient also !

Make sure the weights you are targetting did not change afterwards

Related