Freeze custom model

Hello,
I want to make a custom model that works with the learn.freeze/unfreeze methods.
What I need to implement?
My models has an encoder part that uses:

if freeze:
  with torch.no_grad():
  x = self.encoder(x)
else:
 x = self.encoder(x)

Probably this is not the way , as looking at the freeze_to method on optimizer I can see that only supports an index, and will freeze regarding this value.

1 Like

The freeze/unfreeze is using the underlying parameter groups. So if you want to freeze/unfreeze part of your model, you need to provide a splitter when you instantiate your Learner to tell it which part of your model you want in each group.
There is an example of this in the siameses tutorial.

4 Likes

Merci Sylvain.