使用预训练模型时,可能需要将
(1)预训练好的 backbone 的 参数学习率设置为较小值,
(2)而backbone 之外的部分,需要使用较大的学习率。
python
from collections import OrderedDict
import torch.nn as nn
import torch.optim as optim
net = nn.Sequential(OrderedDict([
("linear1", nn.Linear(10, 20)),
("linear2", nn.Linear(20, 30)),
("linear3", nn.Linear(30, 40))]))
linear3_params = list(map(id, net.linear3.parameters()))
base_params = filter(lambda p: id(p) not in linear3_params, net.parameters())
optimizer = optim.SGD([
{'params': base_params},
{'params': net.linear3.parameters(), 'lr': 0.0005}],
lr=0.001, momentum=0.9)
print(optimizer)
print(optimizer.param_groups[0]['lr'])
print(optimizer.param_groups[1]['lr'])