通过 TensorFlow2 使用学习率计划和学习率预热

Question

我必须使用学习率预热，开始为 CIFAR-10 训练 VGG-19 CNN，并使用学习率预热在前 10000 次迭代（或大约 13 个时期）中将学习率从 0.00001 预热到 0.1。然后，对于剩余的训练，您使用 0.01 的学习率，其中学习率衰减用于在 80 和 120 epoch 时将学习率降低 10 倍。该模型总共需要训练 144 个 epoch。

我使用的是 Python 3 和 TensorFlow2，其中训练数据集有 50000 个示例，批量大小 = 64。一个时期内的训练迭代次数 = 50000/64 = 781 次迭代（大约）。如何在代码中同时使用学习率预热和学习率衰减？

目前，我正在使用学习率衰减：

boundaries = [100000, 110000]
values = [1.0, 0.5, 0.1]

learning_rate_fn = keras.optimizers.schedules.PiecewiseConstantDecay(
    boundaries, values)
print("\nCurrent step value: {0}, LR: {1:.6f}\n".format(optimizer.iterations.numpy(), optimizer.learning_rate(optimizer.iterations)))

但是，我不知道如何使用学习率预热和学习率衰减。

帮助？

Answer 1

使用Transformers 库中的实现怎么样？

from typing import Callable

import tensorflow as tf


class WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):

def __init__(
    self,
    initial_learning_rate: float,
    decay_schedule_fn: Callable,
    warmup_steps: int,
    power: float = 1.0,
    name: str = None,
):
    super().__init__()
    self.initial_learning_rate = initial_learning_rate
    self.warmup_steps = warmup_steps
    self.power = power
    self.decay_schedule_fn = decay_schedule_fn
    self.name = name

def __call__(self, step):
    with tf.name_scope(self.name or "WarmUp") as name:
        # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
        # learning rate will be `global_step/num_warmup_steps * init_lr`.
        global_step_float = tf.cast(step, tf.float32)
        warmup_steps_float = tf.cast(self.warmup_steps, tf.float32)
        warmup_percent_done = global_step_float / warmup_steps_float
        warmup_learning_rate = self.initial_learning_rate * tf.math.pow(warmup_percent_done, self.power)
        return tf.cond(
            global_step_float < warmup_steps_float,
            lambda: warmup_learning_rate,
            lambda: self.decay_schedule_fn(step - self.warmup_steps),
            name=name,
        )

def get_config(self):
    return {
        "initial_learning_rate": self.initial_learning_rate,
        "decay_schedule_fn": self.decay_schedule_fn,
        "warmup_steps": self.warmup_steps,
        "power": self.power,
        "name": self.name,
    }

Answer 2

您可以通过将学习率调度程序设置为lr 参数来将其传递给任何优化程序。例如 -

from tensorlow.keras.optimizers import schedules, RMSProp
boundaries = [100000, 110000]
values = [1.0, 0.5, 0.1]

lr_schedule = schedules.PiecewiseConstantDecay(boundaries, values)
optimizer = keras.optimizers.RMSprop(learning_rate=lr_schedule)