pytorch中为不同model选择不同学习率

我们在训练过程中不免需要为不同的模型中某些模块选择不同的学习率，该如何实现呢？两种实现方式：

• 1 初始化optimizer
• 2 训练过程中为optimzier参数赋值

1 初始化optimizer

只需要在定义optimizer时根据不同model初始化不同的参数。

import torch.nn as nn
import torch.optim


class Model(nn.Module):
    def __init__(self, hidden):
        super().__init__()
        self.fc1 = nn.Linear(hidden, 2 * hidden)
        self.fc2 = nn.Linear(2 * hidden, hidden)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x


model = Model(100)
optimizer = torch.optim.Adam([
    {"params": model.fc1.parameters(), "lr": 0.001},
    {"params": model.fc2.parameters(), "lr": 0.005}
], lr=0.01)
print(optimizer.param_groups)

这里为model中的fc1使用0.001学习率进行优化，对model中的fc2使用0.005学习率进行优化。

2 训练过程中为optimzier参数赋值

for param_group in optimizer.param_groups:
    param_group['lr'] = lr

optimizer.param_groups是一个列表（元素为字典）。