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

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()

次に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)

参考

https://lib.yanxishe.com/document/TensorFlow/api/tf.keras.optimizers.schedules.LearningRateSchedule