Изменение шага learning rate в TensorFlow Keras

Для сходимости к лучшему результату при обучении одним из ключевых параметров является learning_rate, который говорит оптимизатору на каждом шаге сдвинутся на какое значение. Это определяется примерно так:

opt = keras.optimizers.Adam(learning_rate=0.01)

Однако в большинстве случаев постоянное значение learning_rate не является оптимальным. Существуют несколько способов изменять это значение во время обучения.

 Расписание в оптимизаторе

Вместо указания конкретного значения можно подать объект расписания tf.keras.optimizers.schedules.LearningRateSchedule. Например, так:

steps_per_epoch = 200

boundaries = [steps_per_epoch * n for n in [25,50,75]]

values = [1e-3,1e-4,1e-5,1e-6]

lr_sched = optimizers.schedules.PiecewiseConstantDecay(boundaries, values)

optimizer = optimizers.Adam(lr_sched)

Здесь оптимизатор будет изменять значения 3 раза - на 25, 50 и 75 шагах. Начальное значение learning_rate = 1e-3

 Изменение learning_rate, реализованное в виде callback

Второй способ реализовать изменения learning_rate - завести callback. Это можно сделать, например, следующим образом:

def lrfn(current_step, num_warmup_steps, lr_max, num_cycles=0.50, num_training_steps=N_EPOCHS):   

    if current_step < num_warmup_steps:

        if WARMUP_METHOD == 'log':

            return lr_max * 0.10 ** (num_warmup_steps - current_step)

        else:

            return lr_max * 2 ** -(num_warmup_steps - current_step)

    else:

        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))

         return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) * lr_max

LR_SCHEDULE = [lrfn(step, num_warmup_steps=N_WARMUP_EPOCHS, lr_max=LR_MAX, num_cycles=0.50) for step in range(N_EPOCHS)]

lr_callback = tf.keras.callbacks.LearningRateScheduler(lambda step: LR_SCHEDULE[step], verbose=1)

После чего call_back заводится при обучении модели:

history = model.fit(

    x=INPUTS,

    steps_per_epoch=STEPS,

    epochs=EPOCHS,   

    batch_size=BATCH_SIZE,

    validation_data=validation_data,

    callbacks=[

        lr_callback,

        WeightDecayCallback(),

    ],

    verbose = VERBOSE,

)

Конкретно в этом случае значение learning_rate меняется на каждом шаге. Это можно посмотреть по графику:

Эта изображение было получено с помощью следующего кода:

def plot_lr_schedule(lr_schedule, epochs):

    fig = plt.figure(figsize=(20, 10))

    plt.plot([None] + lr_schedule + [None])

    # X Labels

    x = np.arange(1, epochs + 1)

    x_axis_labels = [i if epochs <= 40 or i % 5 == 0 or i == 1 else None for i in range(1, epochs + 1)]

    plt.xlim([1, epochs])

    plt.xticks(x, x_axis_labels) # set tick step to 1 and let x axis start at 1

    # Increase y-limit for better readability

    plt.ylim([0, max(lr_schedule) * 1.1])

    # Title

    schedule_info = f'start: {lr_schedule[0]:.1E}, max: {max(lr_schedule):.1E}, final: {lr_schedule[-1]:.1E}'

    plt.title(f'Step Learning Rate Schedule, {schedule_info}', size=18, pad=12)   

    # Plot Learning Rates

    for x, val in enumerate(lr_schedule):

        if epochs <= 40 or x % 5 == 0 or x is epochs - 1:

            if x < len(lr_schedule) - 1:

                if lr_schedule[x - 1] < val:

                    ha = 'right'

                else:

                    ha = 'left'

            elif x == 0:

                ha = 'right'

            else:

                ha = 'left'

            plt.plot(x + 1, val, 'o', color='black');

            offset_y = (max(lr_schedule) - min(lr_schedule)) * 0.02

            plt.annotate(f'{val:.1E}', xy=(x + 1, val + offset_y), size=12, ha=ha)

   

    plt.xlabel('Epoch', size=16, labelpad=5)

    plt.ylabel('Learning Rate', size=16, labelpad=5)

    plt.grid()

    plt.show()

plot_lr_schedule(LR_SCHEDULE, epochs=N_EPOCHS)