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