Pytorch固定部分参数(只训练部分层)

在迁移学习中我们经常会用到预训练模型,并在预训练模型的基础上添加额外层。训练时先将预训练层参数固定,只训练额外添加的部分。完了之后再全部训练微调。

在pytorch 固定部分参数训练时需要在优化器中施加过滤。

需要自己过滤

optimizer.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3)

另外,如果是Variable,则可以初始化时指定

j = Variable(torch.randn(5,5), requires_grad=True)

但是如果是(神经网络层)

m = nn.Linear(10,10)

是没有requires_grad传入的,m.requires_grad也没有,需要

for i in m.parameters():
    i.requires_grad=False

另外一个小技巧就是在nn.Module里,可以在中间插入这个

for p in self.parameters():
    p.requires_grad=False

这样前面的参数就是False,而后面的不变

 

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)

        for p in self.parameters():
            p.requires_grad=False

        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

 

 1 class RESNET_attention(nn.Module):
 2     def __init__(self, model, pretrained):
 3         super(RESNET_attetnion, self).__init__()
 4         self.resnet = model(pretrained)
 5         for p in self.parameters():
 6             p.requires_grad = False
 7         self.f = nn.Conv2d(2048, 512, 1)
 8         self.g = nn.Conv2d(2048, 512, 1)
 9         self.h = nn.Conv2d(2048, 2048, 1)
10         self.softmax = nn.Softmax(-1)
11         self.gamma = nn.Parameter(torch.FloatTensor([0.0]))
12         self.avgpool = nn.AvgPool2d(7, stride=1)
13         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)


原文:

[1] https://blog.csdn.net/guotong1988/article/details/79739775

[2] https://blog.csdn.net/weixin_34261739/article/details/87555871 

 

posted on 2019-07-17 11:19  那抹阳光1994  阅读(29382)  评论(0编辑  收藏  举报

导航