How can I only update some specific tensors in network with pytorch?

user_7644710 0
Viewed 1557

For instance, I want to only update all cnn weights in Resnet in the first 10 epochs and freeze the others.
And from 11th epoch, I wanna change to update the whole model.
How can I achieve the goal?

conv-neural-network pytorch deep-learning machine-learning image-processing
2 Answers

You can set the learning rate (and some other meta-parameters) per parameters group. You only need to group your parameters according to your needs.
For example, setting different learning rate for conv layers:

import torch
import itertools
from torch import nn

conv_params = itertools.chain.from_iterable([m.parameters() for m in model.children()
                                             if isinstance(m, nn.Conv2d)])
other_params = itertools.chain.from_iterable([m.parameters() for m in model.children()
                                              if not isinstance(m, nn.Conv2d)]) 
optimizer = torch.optim.SGD([{'params': other_params},
                             {'params': conv_params, 'lr': 0}],  # set init lr to 0
                            lr=lr_for_model)

You can later access the optimizer param_groups and modify the learning rate.

See per-parameter options for more information.

user_17144100

very straightforward, as PYTORCH recreate the computational graph on the fly.

for p in resnet.parameters():
    p.requires_grad = False # this will freeze the module from training suppose that resnet is one of your module

if you have multiple modules, simply loop over it. then after 10 epoch, you simply call 

for p in network.parameters():
    p.requires_grad = True # suppose your whole network is the 'network' module
user_72130530
Related