EpochではなくStepごとに学習率制御[Keras]

2月 13, 2023

Transformerで適用されている下式のウォームアップを実装する方法を紹介

https://arxiv.org/pdf/1706.03762.pdf

tf.kerasでは学習率の制御にはtf.keras.callbacks.LearningRateSchedulerを使ってEpoch単位で変更する例が多いが，step毎(つまりバッチ単位)で学習率を制御することができない．
そこで，tf.keras.optimizers.schedules.LearningRateScheduleを用いて新たなクラス(ここではWarmupScheduleという名前にした)をOptimizerのlearning_rateに直接渡すことでステップ毎の学習率制御できるため，これを使って実装する.

先ずはクラスを作成

class WarmupSchedule(keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self):
        self.num_warmup_steps = 4000
        self.d_model = 512

    def __call__(self, step):
        step = step + 1
        step = K.cast(step, dtype="float32")
        d_model = K.cast(self.d_model, dtype="float32")
        num_warmup_steps = K.cast(self.num_warmup_steps, dtype="float32")
        return self.d_model ** -0.5 * K.min([step ** -0.5, step * (self.num_warmup_steps ** -1.5)])

次にOptimizerを定義

lr_schedule = WarmupSchedule()
optimizer = Adam(learning_rate=lr_schedule)

あとはコンパイルして動かすのみ

model.compile(loss="mse", optimizer=optimizer)
history = model.fit(x_train, y_train, validation_data=(x_test, y_test), epochs=epochs, batch_size=batch_size)