Search code examples
pytorch

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

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

Solution

  • According to @PlainRavioli , it's not possible yet and you can set the gradient to zero so the current weights do not change. But you have to do this after calling loss.backward() and before calling optimizer.step().

    So being fc = nn.Linear(n,3), for freezing the parameters of the third output:

    loss.backward()
fc.weight.grad[2,:] = torch.zeros_like(fc.weight.grad[2,:])
fc.bias.grad[2] = torch.zeros_like(fc.bias.grad[2])
optimizer.step()

    Calling loss.backward() computes dloss/dx and do x.grad += dloss/dx. So before this operation, gradients are set to None.