KerasのLearningRateSchedulerとPyTorchのLambdaLRの微妙な違い

---

学習率の調整は大事です。エポック後に学習率を減衰させる際、現在のエポックを引数として更新後の学習率を返す関数を与えると便利なことが多いです。この操作はKeras,PyTorchどちらでもできますが、扱い方が微妙に違うところがあります。ここを知らないでKerasの感覚のままPyTorchでやったらハマりまくったのでメモとして書いておきます。

Kerasの場合は「更新後の学習率」を返す

Kerasの場合はわかりやすいです。エポックを引数として、更新後の学習率をそのまま返す関数を用意すればよいです。以下のコードの場合は「lr_scheduler」という関数ですね。

import keras
from keras import layers
import keras.backend as K

def mnist_mlp():
    input = layers.Input((784,))
    x = layers.Dense(128, activation="relu")(input)
    x = layers.Dense(10, activation="softmax")(x)
    return keras.models.Model(input, x)

# Kerasの場合は実際の学習率を与える
def lr_scheduler(epoch):
    initial_lr = 1e-3
    return (initial_lr - 1e-8) * (10 - epoch) / 10.0

# 確認用コールバック
class PrintCallback(keras.callbacks.Callback):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def on_epoch_end(self, epoch, logs):
        print("Current learning rate = ", K.eval(self.model.optimizer.lr))

# 訓練
def train():
    (X_train, y_train), (X_test, y_test) = keras.datasets.mnist.load_data()
    X_train = (X_train / 255.0).reshape(60000, -1)
    y_train = keras.utils.to_categorical(y_train)

    model = mnist_mlp()
    model.compile(keras.optimizers.Adam(1e-3), "categorical_crossentropy")

    cb = PrintCallback(model)
    lr_scheduling=keras.callbacks.LearningRateScheduler(lr_scheduler)
    model.fit(X_train, y_train, epochs=10, callbacks=[cb, lr_scheduling])

if __name__ == "__main__":
    train()

これはMNIST＋多層パーセプトロンの例です。PrintCallbackというコールバックで現在の学習率を確認しています。確認用なので大した意味はないです。学習率は初期を1e-3として、最終的にほぼ0になるように線形で落としています。10エポック訓練させます。

出力は次のようになります。

60000/60000 [==============================] - 7s 112us/step - loss: 0.2588
Current learning rate =  0.00099999
Epoch 2/10
60000/60000 [==============================] - 4s 60us/step - loss: 0.1132
Current learning rate =  0.000899991
Epoch 3/10
60000/60000 [==============================] - 4s 60us/step - loss: 0.0771
Current learning rate =  0.000799992
Epoch 4/10
60000/60000 [==============================] - 4s 61us/step - loss: 0.0559
Current learning rate =  0.000699993
Epoch 5/10
60000/60000 [==============================] - 4s 62us/step - loss: 0.0428
Current learning rate =  0.000599994
Epoch 6/10
60000/60000 [==============================] - 4s 62us/step - loss: 0.0326
Current learning rate =  0.000499995
Epoch 7/10
60000/60000 [==============================] - 4s 61us/step - loss: 0.0250
Current learning rate =  0.000399996
Epoch 8/10
60000/60000 [==============================] - 4s 61us/step - loss: 0.0199
Current learning rate =  0.000299997
Epoch 9/10
60000/60000 [==============================] - 4s 63us/step - loss: 0.0157
Current learning rate =  0.000199998
Epoch 10/10
60000/60000 [==============================] - 4s 61us/step - loss: 0.0129
Current learning rate =  9.9999e-05

うまくいきました。Kerasの場合は更新後の学習率を返せばいいです。

PyTorchの場合は「ベースの学習率に対する倍率」を返す

似たようなことはPyTorchのLambdaLRを使えばいいですが、こちらはベースの学習率に対する倍率を指定します。Kerasのように更新後の学習率をダイレクト指定ではないです。

import torch
from torch import nn
import torchvision
from torchvision import transforms
import statistics

class MnistMlp(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(            
            nn.Linear(784, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        return self.model(x)

def load_dataset():
    transform = transforms.Compose(
        [transforms.ToTensor()]
    )
    dataset = torchvision.datasets.MNIST(root=".data", train=True, transform=transform, download=True)
    loader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)
    return loader

def lr_scheduling(epoch):
    # ベースの学習率に対する「倍率」なので注意
    return 1.0 *  (10 - epoch) / 10.0

def train():
    train_loader = load_dataset()

    model = MnistMlp().to("cuda")
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss = nn.CrossEntropyLoss()
    scheduler = torch.optim.lr_scheduler.LambdaLR(opt, lr_scheduling)

    for i in range(10):
        print("Epoch", i+1, "/ 10")
        # 現在の学習率を取得
        for param in opt.param_groups:
            current_lr = param["lr"]

        logs = []
        for X, y in train_loader:
            X, y = X.to("cuda"), y.to("cuda")
            X = X.view(X.size(0), -1)

            opt.zero_grad()

            y_pred = model(X)
            current_loss = loss(y_pred, y)
            current_loss.backward()
            opt.step()

            logs.append(current_loss.item())

        # 学習率の更新
        scheduler.step()
        print("Current learning rate =", current_lr, "Loss =", statistics.mean(logs))

if __name__ == "__main__":
    train()

ベースの学習率とはオプティマイザを定義したときに与えたlrの値です。ここでは1e-3です。

lr_schedulingという関数に注目しましょう。ここで初期学習率の値は一切でていません。あくまで比率指定です。なぜそうなのかというとPyTorchのソースを見てみましょう。

LambdaLRの実際に更新後の学習率を取得する部分です。

    def get_lr(self):
        return [base_lr * lmbda(self.last_epoch)
                for lmbda, base_lr in zip(self.lr_lambdas, self.base_lrs)]

「self.base_lrs」には学習率更新の関数のリストが記録されています。やはり、ベースの学習率に対する比ですよね。

実際最初に出したコードの出力を確認すると次のようになります。

Epoch 1 / 10
Current learning rate = 0.001 Loss = 0.29384363802075386
Epoch 2 / 10
Current learning rate = 0.0009000000000000001 Loss = 0.12908256237705548
Epoch 3 / 10
Current learning rate = 0.0008 Loss = 0.08823229066530863
Epoch 4 / 10
Current learning rate = 0.0007 Loss = 0.06536479882498582
Epoch 5 / 10
Current learning rate = 0.0006 Loss = 0.050669593796134
Epoch 6 / 10
Current learning rate = 0.0005 Loss = 0.04010948433230321
Epoch 7 / 10
Current learning rate = 0.0004 Loss = 0.031612708161771294
Epoch 8 / 10
Current learning rate = 0.0003 Loss = 0.025767304584383966
Epoch 9 / 10
Current learning rate = 0.0002 Loss = 0.021104651134461163
Epoch 10 / 10
Current learning rate = 0.0001 Loss = 0.017641670887172223

正しく学習率が更新されているのが確認できます。

LambdaLRでKerasの感覚で実際の学習率を返してしまうと、学習率が2乗されてしまい、学習率がものすごく低くなってしまいます。あたかも勾配が消失しているように見えますが、これは勾配消失ではなく、ただ単に学習率のスケジューリングが正常に動作していないというだけです。ここ知らなくてハマりました。

まとめ

• Kerasの場合はLearningRateSchedulerで、実際の学習率を返す
• PyTorchの場合はLambdaLRで、ベースの学習率に対する比率を返す

---

---

Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内

---