Calentamiento de la tasa de aprendizaje con decaimiento del coseno en Keras/TensorFlow

La tasa de aprendizaje es un hiperparámetro importante en las redes de aprendizaje profundo, y dicta directamente el la licenciatura a los que se realizan actualizaciones de pesos, que se estiman para minimizar alguna función de pérdida dada. en EUR:

$$
peso_{t+1} = peso_t – lr * frac{derror}{dpeso_t}
$$

Con una tasa de aprendizaje de 0, el peso actualizado vuelve a ser el mismo: peso_t. La tasa de aprendizaje es efectivamente una perilla que podemos girar para habilitar o deshabilitar el aprendizaje, y tiene una gran influencia sobre la cantidad de aprendizaje que se está produciendo, al controlar directamente el grado de actualización del peso.

Los diferentes optimizadores utilizan las tasas de aprendizaje de manera diferente, pero el concepto subyacente sigue siendo el mismo. No hace falta decir que las tasas de aprendizaje han sido objeto de muchos estudios, artículos y puntos de referencia de los profesionales.

En términos generales, casi todo el mundo está de acuerdo en que una tasa de aprendizaje estática no es suficiente, y se produce algún tipo de reducción de la tasa de aprendizaje en la mayoría de las técnicas que ajustan la tasa de aprendizaje durante el entrenamiento, ya sea monótono, coseno, triangular u otros tipos de reducción.

Una técnica que en los últimos años ha ido afianzándose es calentamiento de la tasa de aprendizaje, que se puede combinar con prácticamente cualquier otra técnica de reducción.

Calentamiento de la tasa de aprendizaje

La idea detrás del calentamiento de la tasa de aprendizaje es simple. En las primeras etapas del entrenamiento, los pesos están lejos de su estado ideal. Esto significa grandes actualizaciones en todos los ámbitos, que pueden verse como "correcciones excesivas" para cada peso, donde la actualización drástica de otro puede negar la actualización de algún otro peso, lo que hace que las etapas iniciales del entrenamiento sean más inestables.

Estos cambios se solucionan, pero se pueden evitar teniendo una tasa de aprendizaje pequeña para empezar, alcanzando un estado subóptimo más estable y luego aplicando una tasa de aprendizaje mayor. Puede facilitar las actualizaciones de la red, en lugar de golpearlas.

¡Ese es el calentamiento de la tasa de aprendizaje! Comenzando con una tasa de aprendizaje baja (o 0) y aumentando a una tasa de aprendizaje inicial (con lo que comenzaría de todos modos). Este incremento puede seguir realmente cualquier función, pero es comúnmente lineal.

Después de alcanzar la tasa inicial, se pueden aplicar otros programas como la disminución del coseno, la reducción lineal, etc. para reducir progresivamente la tasa hasta el final del entrenamiento. El calentamiento de la tasa de aprendizaje suele ser parte de un programa de dos programas, donde el calentamiento de LR es el primero, mientras que otro programa se hace cargo después de que la tasa ha alcanzado un punto de inicio.

En esta guía, implementaremos un calentamiento de la tasa de aprendizaje en Keras/TensorFlow como keras.optimizers.schedules.LearningRateSchedule subclase y keras.callbacks.Callback llamar de vuelta. La tasa de aprendizaje aumentará de 0 a target_lr y aplique el decaimiento del coseno, ya que este es un programa secundario muy común. Como de costumbre, Keras simplifica la implementación de soluciones flexibles de varias maneras y las envía con su red.

Tasa de aprendizaje con devoluciones de llamada de Keras

La forma más sencilla de implementar cualquier programa de tasa de aprendizaje es mediante la creación de una función que tome la lr parámetro (float32), lo pasa a través de alguna transformación y lo devuelve. Esta función luego se pasa al LearningRateScheduler devolución de llamada, que aplica la función a la tasa de aprendizaje.

Ahora, la tf.keras.callbacks.LearningRateScheduler() pasa el número de época a la función que usa para calcular la tasa de aprendizaje, que es bastante tosca. LR Warmup debe hacerse en cada paso (lote), no época, por lo que tendremos que derivar un global_step (en todas las épocas) para calcular la tasa de aprendizaje en su lugar, y subclasificar el Callback class para crear una devolución de llamada personalizada en lugar de simplemente pasar la función, ya que necesitaremos pasar argumentos en cada llamada, lo cual es imposible cuando solo se pasa la función:

def func():
    return ...
    
keras.callbacks.LearningRateScheduler(func)

Este enfoque es favorable cuando no desea un alto nivel de personalización y no desea interferir con la forma en que Keras trata el lr, y especialmente si desea utilizar devoluciones de llamada como ReduceLROnPlateau() ya que solo puede funcionar con un flotante lr. Implementemos un calentamiento de la tasa de aprendizaje mediante una devolución de llamada de Keras, comenzando con una función de conveniencia:

def lr_warmup_cosine_decay(global_step,
                           warmup_steps,
                           hold = 0,
                           total_steps=0,
                           start_lr=0.0,
                           target_lr=1e-3):
    
    learning_rate = 0.5 * target_lr * (1 + np.cos(np.pi * (global_step - warmup_steps - hold) / float(total_steps - warmup_steps - hold)))

    
    warmup_lr = target_lr * (global_step / warmup_steps)

    
    
    if hold > 0:
        learning_rate = np.where(global_step > warmup_steps + hold,
                                 learning_rate, target_lr)
    
    learning_rate = np.where(global_step < warmup_steps, warmup_lr, learning_rate)
    return learning_rate

En cada paso, calculamos la tasa de aprendizaje y la tasa de aprendizaje de calentamiento (ambos elementos del cronograma), con respecto a la start_lr y target_lr. start_lr por lo general comenzará en 0.0, Mientras que el target_lr depende de su red y optimizador – 1e-3 podría no ser un buen valor predeterminado, así que asegúrese de establecer su objetivo de inicio LR cuando llame al método.

Si global_step en el entrenamiento es más alto que el warmup_steps hemos establecido: usamos el programa de decaimiento del coseno LR. Si no, significa que todavía estamos calentando, por lo que se usa el LR de calentamiento. Si el hold se establece el argumento, mantendremos el target_lr para ese número de pasos después del calentamiento y antes de que decaiga el coseno. np.where() proporciona una gran sintaxis para esto:

np.where(condition, value_if_true, value_if_false)

Puedes visualizar la función con:

steps = np.arange(0, 1000, 1)
lrs = []

for step in steps:
  lrs.append(lr_warmup_cosine_decay(step, total_steps=len(steps), warmup_steps=100, hold=10))
plt.plot(lrs)

Ahora, querremos usar esta función como parte de una devolución de llamada y pasar el paso del optimizador como el global_step en lugar de un elemento de una matriz arbitraria, o puede realizar el cálculo dentro de la clase. Vamos a subclasificar el Callback clase:

from keras import backend as K

class WarmupCosineDecay(keras.callbacks.Callback):
    def __init__(self, total_steps=0, warmup_steps=0, start_lr=0.0, target_lr=1e-3, hold=0):

        super(WarmupCosineDecay, self).__init__()
        self.start_lr = start_lr
        self.hold = hold
        self.total_steps = total_steps
        self.global_step = 0
        self.target_lr = target_lr
        self.warmup_steps = warmup_steps
        self.lrs = []

    def on_batch_end(self, batch, logs=None):
        self.global_step = self.global_step + 1
        lr = model.optimizer.lr.numpy()
        self.lrs.append(lr)

    def on_batch_begin(self, batch, logs=None):
        lr = lr_warmup_cosine_decay(global_step=self.global_step,
                                    total_steps=self.total_steps,
                                    warmup_steps=self.warmup_steps,
                                    start_lr=self.start_lr,
                                    target_lr=self.target_lr,
                                    hold=self.hold)
        K.set_value(self.model.optimizer.lr, lr)

Primero, definimos el constructor de la clase y realizamos un seguimiento de sus campos. En cada lote que termine, aumentaremos el paso global, tomaremos nota del LR actual y lo agregaremos a la lista de LR hasta el momento. Al comienzo de cada lote, calcularemos el LR usando el lr_warmup_cosine_decay() y configure ese LR como el LR actual del optimizador. Esto se hace con el backend set_value().

Una vez hecho esto, simplemente calcule los pasos totales (longitud/tamaño del lote*épocas) y tome una parte de ese número para su warmup_steps:

total_steps = len(train_set)*config['EPOCHS']



warmup_steps = int(0.05*total_steps)

callback = WarmupCosineDecay(total_steps=total_steps, 
                             warmup_steps=warmup_steps,
                             hold=int(warmup_steps/2), 
                             start_lr=0.0, 
                             target_lr=1e-3)

Finalmente, construya su modelo y proporcione la devolución de llamada en el fit() llamada:

model = keras.applications.EfficientNetV2B0(weights=None, 
                                            classes=n_classes, 
                                            input_shape=[224, 224, 3])
  
model.compile(loss="sparse_categorical_crossentropy",
                  optimizer='adam',
                  jit_compile=True,
                  metrics=['accuracy'])

Al final del entrenamiento, puede obtener y visualizar los LR modificados a través de:

lrs = callback.lrs 
plt.plot(lrs)

Si traza el historial de un modelo entrenado con y sin calentamiento LR, verá una clara diferencia en la estabilidad del entrenamiento:

Tasa de aprendizaje con la subclase LearningRateSchedule

Una alternativa a la creación de una devolución de llamada es crear un LearningRateSchedule subclase, que no manipula el LR, lo reemplaza. Este enfoque le permite profundizar un poco más en el backend de Keras/TensorFlow, pero cuando se usa, no se puede combinar con otras devoluciones de llamada relacionadas con LR, como ReduceLROnPlateau(), que trata los LR como números de coma flotante.

Además, el uso de la subclase requerirá que la haga serializable (sobrecarga get_config()) a medida que se convierte en parte del modelo, si desea guardar los pesos del modelo. Otra cosa a tener en cuenta es que la clase esperará trabajar exclusivamente con tf.Tensors. Afortunadamente, la única diferencia en la forma en que trabajamos será llamar tf.func() en lugar de np.func() ya que las API de TensorFlow y NumPy son increíblemente similares y compatibles.

Reescribamos la conveniencia lr_warmup_cosine_decay() para usar las operaciones de TensorFlow en su lugar:

def lr_warmup_cosine_decay(global_step,
                           warmup_steps,
                           hold = 0,
                           total_steps=0,
                           start_lr=0.0,
                           target_lr=1e-3):
    
    
    learning_rate = 0.5 * target_lr * (1 + tf.cos(tf.constant(np.pi) * (global_step - warmup_steps - hold) / float(total_steps - warmup_steps - hold)))

    
    warmup_lr = target_lr * (global_step / warmup_steps)

    
    
    if hold > 0:
        learning_rate = tf.where(global_step > warmup_steps + hold,
                                 learning_rate, target_lr)
    
    learning_rate = tf.where(global_step < warmup_steps, warmup_lr, learning_rate)
    return learning_rate

Con la función de conveniencia, podemos subclasificar el LearningRateSchedule clase. En cada __call__() (lote), calcularemos el LR usando la función y lo devolveremos. Naturalmente, también puede empaquetar el cálculo dentro de la clase subclase.

La sintaxis es más limpia que la Callback sublcass, principalmente porque tenemos acceso a la step campo, en lugar de hacer un seguimiento de él por nuestra cuenta, sino que también hace que sea un poco más difícil trabajar con propiedades de clase, en particular, hace que sea difícil extraer el lr de un tf.Tensor() en cualquier otro tipo para realizar un seguimiento en una lista. Esto se puede eludir técnicamente ejecutando en modo ansioso, pero presenta una molestia para realizar un seguimiento del LR con fines de depuración y es mejor evitarlo:

class WarmUpCosineDecay(keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, start_lr, target_lr, warmup_steps, total_steps, hold):
        super().__init__()
        self.start_lr = start_lr
        self.target_lr = target_lr
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.hold = hold

    def __call__(self, step):
        lr = lr_warmup_cosine_decay(global_step=step,
                                    total_steps=self.total_steps,
                                    warmup_steps=self.warmup_steps,
                                    start_lr=self.start_lr,
                                    target_lr=self.target_lr,
                                    hold=self.hold)

        return tf.where(
            step > self.total_steps, 0.0, lr, name="learning_rate"
        )

Los parámetros son los mismos y se pueden calcular de la misma manera que antes:

total_steps = len(train_set)*config['EPOCHS']



warmup_steps = int(0.05*total_steps)

schedule = WarmUpCosineDecay(start_lr=0.0, target_lr=1e-3, warmup_steps=warmup_steps, total_steps=total_steps, hold=warmup_steps)

Y la canalización de entrenamiento solo difiere en que configuramos el LR del optimizador en el schedule:

model = keras.applications.EfficientNetV2B0(weights=None, 
                                            classes=n_classes, 
                                            input_shape=[224, 224, 3])
  
model.compile(loss="sparse_categorical_crossentropy",
                  optimizer=tf.keras.optimizers.Adam(learning_rate=schedule),
                  jit_compile=True,
                  metrics=['accuracy'])

history3 = model.fit(train_set,
                    epochs = config['EPOCHS'],
                    validation_data=valid_set)

Si desea guardar el modelo, el WarmupCosineDecay el horario tendrá que anular el get_config() método:

    def get_config(self):
        config = {
          'start_lr': self.start_lr,
          'target_lr': self.target_lr,
          'warmup_steps': self.warmup_steps,
          'total_steps': self.total_steps,
          'hold': self.hold
        }
        return config

Finalmente, al cargar el modelo, tendrás que pasar un WarmupCosineDecay como un objeto personalizado:

model = keras.models.load_model('weights.h5', 
                                custom_objects={'WarmupCosineDecay', WarmupCosineDecay})

Conclusión

En esta guía, echamos un vistazo a la intuición detrás del calentamiento de la tasa de aprendizaje, una técnica común para manipular la tasa de aprendizaje mientras se entrenan redes neuronales.

Hemos implementado un calentamiento de tasa de aprendizaje con decaimiento de coseno, el tipo más común de reducción de LR combinado con calentamiento. Puede implementar cualquier otra función para la reducción, o no reducir la tasa de aprendizaje en absoluto, dejándolo para otras devoluciones de llamada como ReduceLROnPlateau(). Implementamos el calentamiento de la tasa de aprendizaje como una devolución de llamada de Keras, así como un programa de optimización de Keras y trazamos la tasa de aprendizaje a través de las épocas.