CODEWITHSUNDEEP
machine-learningpytorchlayer
Tclha
Tclha

Reputation: 105

Unfreeze model Layer by Layer in PyTorch

I'm working with a PyTorch model from here (T2T_ViT_7).

I'm trying to freeze all layers except the last (head) layer and train it on my dataset. I want to evaluate its performance, and then unfreeze layers one by one and train it each time and see how it performs.

To initially freeze all the layers and and just unfreeze the head layer, I used:

for param in model.parameters():
    param.requires_grad_(False)

model.head.requires_grad_(True)

Now I want to start from the bottom, and start unfreezing layers one by one. How can I do this? Do I use model.modules() or maybe model.children()?

Thank you!

Upvotes: 0

Views: 1543

Answers (1)

Ivan
Ivan

Reputation: 40708

If by layers you mean each block inside of model.blocks, then you can use nn.Module.children (// nn.Module.named_children). This will return all direct submodules, while the nn.Module.modules returns all submodules recursively.

Since model.blocks is a nn.ModuleList, you can slice the blocks to only select the last n layers. Something like that:

model.blocks[-n:].requires_grad_(False)

Upvotes: 2

Related Questions