pytorch中optimizer为不同参数设置不同的学习率

在pytorch中已经实现了一些常见的优化器,例如Adam、SGD、Adagrad、RMsprop等,但是有些任务中我们需要设定不同的学习策略,例如给模型的不同参数设置不同的学习率

class Linear(nn.Module):
    def __init__(self):
        super().__init__()
        self.w1 = nn.Parameter(torch.randn(3, 4))
        self.b1 = nn.Parameter(torch.randn(1, 3))
        
        self.w2 = nn.Parameter(torch.randn(3, 2))
        self.b2 = nn.Parameter(torch.randn(1, 2))
        
    def forward(self, x):
        x = F.linear(x, self.w1, self.b1)
        return F.linear(x, self.w2, self.b2)

该网络我们定义了4个可学习参数,2个是权重矩阵w,2个是偏置矩阵b,我们假定要为w矩阵设置学习率为1e-2,而为b矩阵设置为1e-3。

实现这种需求其实很简单,只需要在定义优化器时传入一个字典,分别传入需要优化的参数列表,以及对应的学习率。

model = Linear()

w_params = [param for name, param in model.named_parameters() if 'w' in name]
b_params = [param for name, param in model.named_parameters() if 'b' in name]

optimizer = torch.optim.Adam([
    {'params': w_params, 'lr': 1e-2},
    {'params': b_params, 'lr': 1e-3}
])

你可能感兴趣的:(PyTorch,pytorch,深度学习,人工智能,python,神经网络)