pytorch 关于权重(初始化、固定)

权重初始化:

emb那种可以直接from_pretrained

 

固定权重

利用requires_grad=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)

然后在optimizer里需要过滤

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

或者可以def不同的func,return需要更新的param,然后在优化器里用,这样可以让网络不同的时刻(epoch),可以用不同的优化器进行更新参数。比如刚开始更新全部参数,训练后期就只更新上层参数

#在model的class里
def XXoptim(self, lr):
        return [
            {'params': self.XX.parameters(), 'lr': lr},
        ]

#在优化器里
XXoptimizer = torch.optim.SGD(model.XXoptim(lr=0.001), args.MOMENTUM, args.WEIGHT_DECAY)

 

你可能感兴趣的:(深度学习框架)