Learning rate scheduling en TensorFlow

¿Por qué programar la tasa de aprendizaje?

La tasa de aprendizaje (learning rate) controla el tamaño de los pasos durante el descenso del gradiente. Un valor fijo durante todo el entrenamiento suele ser subóptimo:

• Un lr alto al principio permite explorar el espacio de parámetros rápidamente.
• Un lr alto al final puede hacer que el modelo oscile alrededor del mínimo sin converger.
• Un lr progresivamente más pequeño permite un ajuste fino en las etapas tardías del entrenamiento.

Los schedulers de aprendizaje automatizan este proceso.

ExponentialDecay: decaimiento exponencial

Reduce el learning rate de forma exponencial según el número de pasos:

lr = lr_inicial × decay_rate ^ (paso / decay_steps)

import tensorflow as tf
import numpy as np

# Configurar decaimiento exponencial
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate=0.01,
    decay_steps=1000,      # cada 1000 pasos
    decay_rate=0.9,        # se multiplica por 0.9
    staircase=False        # True = decaimiento escalonado; False = continuo
)

# Visualizar cómo evoluciona el lr
pasos = np.arange(0, 5000, 100)
lr_valores = [lr_schedule(paso).numpy() for paso in pasos]
print(f"LR en paso 0:    {lr_valores[0]:.6f}")
print(f"LR en paso 1000: {lr_schedule(1000).numpy():.6f}")
print(f"LR en paso 5000: {lr_schedule(5000).numpy():.6f}")

# Usar con un optimizador
optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
print(f"Optimizador configurado con ExponentialDecay.")

CosineDecay y CosineDecayRestarts

CosineDecay sigue una curva de coseno desde lr_inicial hasta lr_mínimo:

import tensorflow as tf

# Cosine decay estándar
cosine_decay = tf.keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=0.01,
    decay_steps=5000,
    alpha=0.0001           # lr mínimo = alpha * lr_inicial
)

print(f"LR en paso 0:    {cosine_decay(0).numpy():.6f}")
print(f"LR en paso 2500: {cosine_decay(2500).numpy():.6f}")
print(f"LR en paso 5000: {cosine_decay(5000).numpy():.6f}")

# Cosine Decay con reinicios (SGDR: Stochastic Gradient Descent with Warm Restarts)
cosine_restarts = tf.keras.optimizers.schedules.CosineDecayRestarts(
    initial_learning_rate=0.01,
    first_decay_steps=1000,  # pasos del primer ciclo
    t_mul=2.0,               # cada ciclo dura t_mul × el anterior
    m_mul=0.9,               # lr_inicial del siguiente ciclo = m_mul × anterior
    alpha=0.001
)

print(f"\nCosineDecayRestarts en paso 500:  {cosine_restarts(500).numpy():.6f}")
print(f"CosineDecayRestarts en paso 1000: {cosine_restarts(1000).numpy():.6f}")
print(f"CosineDecayRestarts en paso 1500: {cosine_restarts(1500).numpy():.6f}")

PolynomialDecay: decaimiento polinomial

import tensorflow as tf

poly_decay = tf.keras.optimizers.schedules.PolynomialDecay(
    initial_learning_rate=0.01,
    decay_steps=5000,
    end_learning_rate=0.0001,
    power=2.0,             # grado del polinomio (1.0 = lineal)
    cycle=False
)

print(f"LR inicial:    {poly_decay(0).numpy():.6f}")
print(f"LR en paso 5000: {poly_decay(5000).numpy():.6f}")

ReduceLROnPlateau: reducción adaptativa como callback

Este callback reduce el lr cuando una métrica deja de mejorar. Es muy útil cuando no se conoce de antemano el número de pasos:

import tensorflow as tf
import numpy as np

# Modelo sencillo de ejemplo
modelo = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu', input_shape=(50,)),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
modelo.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# Callback ReduceLROnPlateau
reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(
    monitor='val_loss',     # métrica a observar
    factor=0.5,             # nuevo lr = lr × factor
    patience=3,             # épocas sin mejora antes de reducir
    min_lr=1e-6,            # lr mínimo permitido
    min_delta=0.001,        # mejora mínima para contar como mejora
    cooldown=1,             # épocas de espera tras una reducción
    verbose=1
)

# Datos de ejemplo
X = np.random.randn(1000, 50).astype('float32')
y = np.random.randint(0, 10, 1000).astype('int32')

historial = modelo.fit(
    X, y,
    epochs=30,
    validation_split=0.2,
    callbacks=[reduce_lr],
    verbose=1
)

print(f"\nLR final: {modelo.optimizer.learning_rate.numpy():.8f}")

Learning rate con warm-up personalizado

El warm-up es una técnica donde el lr comienza muy bajo y sube gradualmente durante las primeras épocas, evitando inestabilidades al inicio:

import tensorflow as tf
import numpy as np

class WarmUpCosineDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
    """Learning rate schedule con warm-up lineal + cosine decay."""

    def __init__(self, target_lr, warmup_steps, total_steps):
        super().__init__()
        self.target_lr = target_lr
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps

    def __call__(self, step):
        step = tf.cast(step, tf.float32)
        warmup_steps = tf.cast(self.warmup_steps, tf.float32)
        total_steps = tf.cast(self.total_steps, tf.float32)

        # Fase de warm-up: lr crece linealmente desde 0 hasta target_lr
        warmup_lr = self.target_lr * step / warmup_steps

        # Fase de cosine decay: lr baja con coseno
        cosine_decay = 0.5 * self.target_lr * (
            1 + tf.cos(np.pi * (step - warmup_steps) / (total_steps - warmup_steps))
        )

        return tf.where(step < warmup_steps, warmup_lr, cosine_decay)

    def get_config(self):
        return {
            'target_lr': self.target_lr,
            'warmup_steps': self.warmup_steps,
            'total_steps': self.total_steps
        }

# Uso
EPOCHS = 20
STEPS_PER_EPOCH = 100
TOTAL_STEPS = EPOCHS * STEPS_PER_EPOCH
WARMUP_STEPS = 2 * STEPS_PER_EPOCH

lr_schedule = WarmUpCosineDecay(
    target_lr=0.001,
    warmup_steps=WARMUP_STEPS,
    total_steps=TOTAL_STEPS
)

optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)

# Verificar
for paso in [0, 100, 200, 1000, 2000]:
    print(f"LR en paso {paso}: {lr_schedule(paso).numpy():.6f}")

Callback LearningRateScheduler

Para usar una función Python arbitraria como schedule:

import tensorflow as tf
import numpy as np

def schedule_por_epoca(epoca, lr_actual):
    """Reduce lr a la mitad cada 5 épocas."""
    if epoca > 0 and epoca % 5 == 0:
        nuevo_lr = lr_actual * 0.5
        print(f"\nÉpoca {epoca}: lr {lr_actual:.6f} → {nuevo_lr:.6f}")
        return nuevo_lr
    return lr_actual

lr_callback = tf.keras.callbacks.LearningRateScheduler(
    schedule=schedule_por_epoca,
    verbose=0
)

modelo = tf.keras.Sequential([
    tf.keras.layers.Dense(32, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(1)
])
modelo.compile(optimizer='adam', loss='mse')

X = np.random.randn(500, 10).astype('float32')
y = np.random.randn(500).astype('float32')

historial = modelo.fit(X, y, epochs=15, callbacks=[lr_callback], verbose=0)
print(f"LR al final: {float(modelo.optimizer.learning_rate):.6f}")

Visualizar la evolución del learning rate

import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

# Comparar schedulers
schedules = {
    'ExponentialDecay': tf.keras.optimizers.schedules.ExponentialDecay(0.01, 1000, 0.9),
    'CosineDecay': tf.keras.optimizers.schedules.CosineDecay(0.01, 5000),
    'PolynomialDecay': tf.keras.optimizers.schedules.PolynomialDecay(0.01, 5000, 0.0001, power=2)
}

pasos = np.arange(0, 5001, 50)

plt.figure(figsize=(10, 5))
for nombre, schedule in schedules.items():
    lr_vals = [float(schedule(p)) for p in pasos]
    plt.plot(pasos, lr_vals, label=nombre)

plt.xlabel('Pasos de entrenamiento')
plt.ylabel('Learning rate')
plt.title('Comparativa de learning rate schedules')
plt.legend()
plt.yscale('log')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()