PyTorchでレイヤーごとに違う学習率を設定する

PyTorchで各レイヤーごとに違うLearning Rateを設定する方法．

例として，以下のようなネットワークを想定する．

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 256)
        self.fc3 = nn.Linear(256, 10)
        
    def forward(self, x):
        h = F.relu(self.fc1(x))
        h = F.relu(self.fc2(h))
        h = F.relu(self.fc3(h))
        
        return h

オプティマイザを定義する際に，第一引数に以下のようにパラメータを渡せば良い．

net = Net()
optimizer = torch.optim.Adam(params=[
    {"params": net.fc1.parameters(), "lr": 1e-5},
    {"params": net.fc2.parameters(), "lr": 1e-4},
    {"params": net.fc3.parameters(), "lr": 1e-3},
])

Adam (
Parameter Group 0
    amsgrad: False
    betas: (0.9, 0.999)
    eps: 1e-08
    lr: 1e-05
    weight_decay: 0

Parameter Group 1
    amsgrad: False
    betas: (0.9, 0.999)
    eps: 1e-08
    lr: 0.0001
    weight_decay: 0

Parameter Group 2
    amsgrad: False
    betas: (0.9, 0.999)
    eps: 1e-08
    lr: 0.001
    weight_decay: 0
)

このようなパラメータの渡し方をすることで，レイヤーごとに違う学習率を設定できる．

ちなみに，この方法を用いる場合，明示的にパラメータを指定しなかったレイヤーはoptimizer.step()の際にパラメータ更新が行われないことに注意． つまり，レイヤーごとに違うパラメータを指定したい場合は，すべてのレイヤーについて明示的にパラメータを設定する必要がある．