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}
])