参考:https://www.cnblogs.com/jiangkejie/p/11199847.html
在迁移学习中我们经常会用到预训练模型,并在预训练模型的基础上添加额外层。训练时先将预训练层参数固定,只训练额外添加的部分。完了之后再全部训练微调。
在pytorch 固定部分参数训练时需要在优化器中施加过滤。
class RESNET_attention(nn.Module):
def __init__(self, model, pretrained):
super(RESNET_attetnion, self).__init__()
self.resnet = model(pretrained)
for p in self.parameters():
p.requires_grad = False
self.f = nn.Conv2d(2048, 512, 1)
self.g = nn.Conv2d(2048, 512, 1)
self.h = nn.Conv2d(2048, 2048, 1)
self.softmax = nn.Softmax(-1)
self.gamma = nn.Parameter(torch.FloatTensor([0.0]))
self.avgpool = nn.AvgPool2d(7, stride=1)
self.resnet.fc = nn.Linear(2048, 10)
note:以上代码复现SAGAN的Attention部分,这不是主要问题
这样就将for循环以上的参数固定, 只训练下面的参数(f,g,h,gamma,fc,等), 但是注意需要在optimizer中添加上这样的一句话filter(lambda p: p.requires_grad, model.parameters()
添加的位置为:
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.0001, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-5)