How to freeze the top layers?

I have code that looks like this

y = model1(x)
z = model2(y)

I would like the optimizer to update the parameters of model1 only. How can I do that?

You can just pass the parameters of model1 to the optimizer:

optimizer = optim.SGD(model1.parameters(), lr=1e-3)

You could also pass both sets of parameters to the optimizer but have the learning rate of one set be zero. Then you could change the learning rate later on.

optimizer = SGD([{"params": model1.parameters()},
                 {"params": model2.parameters(), "lr": 0.}], lr=0.1, momentum=0.9, nesterov=False)

Here I’ve used a default learning rate for the optimizer which would be used with model1 and explicitly set the learning rate for model2. If you had more models or more sets of parameters you could extend this.

1 Like