自定义学习率调度器 TF2 和 Keras

Question

I am trying to write custom learning rate scheduler: cosine annealing with warm-up. But I can't use it neither in Keras, nor in Tensorflow. Below is the code:

import tensorflow as tf
import numpy as np


def make_linear_lr(min_lr, max_lr, number_of_steps):
    def gen_lr(step):
        return (max_lr - min_lr) / number_of_steps * step + min_lr
    return gen_lr

def make_cosine_anneal_lr(learning_rate, alpha, decay_steps):
    def gen_lr(global_step):
        global_step = min(global_step, decay_steps)
        cosine_decay = 0.5 * (1 + np.cos(np.pi * global_step / decay_steps))
        decayed = (1 - alpha) * cosine_decay + alpha
        decayed_learning_rate = learning_rate * decayed
        return decayed_learning_rate
    return gen_lr

def make_cosine_annealing_with_warmup(min_lr, max_lr, number_of_steps, alpha, decay_steps):
    gen_lr_1 = make_linear_lr(min_lr, max_lr, number_of_steps)
    gen_lr_2 = make_cosine_anneal_lr(max_lr, alpha, decay_steps)
    def gen_lr(global_step):
        if global_step < number_of_steps:
            return gen_lr_1(global_step)
        else:
            return gen_lr_2(global_step - number_of_steps)
        
    return gen_lr

class CosineAnnealingWithWarmUP(tf.keras.optimizers.schedules.LearningRateSchedule):
  def __init__(self, min_lr, max_lr, number_of_steps, alpha, decay_steps):
    super(CosineAnnealingWithWarmUP, self).__init__()
    self.gen_lr_ca =  make_cosine_annealing_with_warmup(min_lr, max_lr, number_of_steps, alpha, decay_steps)
  def __call__(self, step):
    return tf.cast(self.gen_lr_ca(step), tf.float32)

learning_rate_fn = CosineAnnealingWithWarmUP(.0000001, 0.01, 10_000, 0, 150_000)
optimizer=tf.keras.optimizers.SGD(
    learning_rate=learning_rate_fn, 
    momentum=0.95)

I use this function in TensorFlow to train my model:

def get_model_train_step_function(model, optimizer, vars_to_fine_tune, batch_size):
  @tf.function
  def train_step_fn(image_tensors,
                    groundtruth_boxes_list,
                    groundtruth_classes_list):
    shapes = tf.constant(batch_size * [[640, 640, 3]], dtype=tf.int32)
    model.provide_groundtruth(
        groundtruth_boxes_list=groundtruth_boxes_list,
        groundtruth_classes_list=groundtruth_classes_list)
    with tf.GradientTape() as tape:
      preprocessed_images = tf.concat(
          [model.preprocess(
              image_tensor
           )[0]
           for image_tensor in image_tensors], axis=0)
      prediction_dict = model.predict(preprocessed_images, shapes)
      losses_dict = model.loss(prediction_dict, shapes)
      total_loss = losses_dict['Loss/localization_loss'] + losses_dict['Loss/classification_loss']
      gradients = tape.gradient(total_loss, vars_to_fine_tune)
      optimizer.apply_gradients(zip(gradients, vars_to_fine_tune))
    return total_loss

  return train_step_fn 

When I try to use it with TensorFlow, passing optimizer in get_model_train_step_function — it works if I remove @tf.function decorator. But it doesn't work with @tf.function, the error says: OperatorNotAllowedInGraphError: using a tf.Tensor as a Python bool is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

How should I write my custom learning rate scheduler? Also, I would like to use this Schedule with Keras. But it doesn't work there at all.

Answer 1

You need to exclude numpy calls and replace python conditionals ("if", "min") by tensorflow operators:

def make_cosine_anneal_lr(learning_rate, alpha, decay_steps):
    def gen_lr(global_step):

        #global_step = min(global_step, decay_steps)

        global_step = tf.minimum(global_step, decay_steps)
        cosine_decay = 0.5 * (1 + tf.cos(3.1415926 * global_step / decay_steps)) # changed np.pi to 3.14
        decayed = (1 - alpha) * cosine_decay + alpha
        decayed_learning_rate = learning_rate * decayed
        return decayed_learning_rate
    return gen_lr

def make_cosine_annealing_with_warmup(min_lr, max_lr, number_of_steps, alpha, decay_steps):
    gen_lr_1 = make_linear_lr(min_lr, max_lr, number_of_steps)
    gen_lr_2 = make_cosine_anneal_lr(max_lr, alpha, decay_steps)
    def gen_lr(global_step):

      #if global_step < number_of_steps:
      #    return gen_lr_1(global_step)
      #else:
      #    return gen_lr_2(global_step - number_of_steps)

      a = global_step < number_of_steps
      a = tf.cast(a, tf.float32)
      b = 1. - a
      return a * gen_lr_1(global_step) + b * gen_lr_2(global_step - number_of_steps)
        
    return gen_lr

Such schedule works from Keras.