How to set a learning rate scheduler in Keras

An optimizer can take broad steps early in training and smaller steps after the model begins to settle. An epoch-level learning rate schedule makes that change repeatable inside model.fit() instead of relying on manual restarts or optimizer edits between runs.

The keras.callbacks.LearningRateScheduler callback invokes a scheduling function at the beginning of every epoch. The function receives a zero-based epoch index and the optimizer's current learning rate, then returns the value Keras should use for that epoch.

A five-epoch run on the JAX backend starts at 0.1 and applies three successive reductions. TensorFlow and PyTorch projects can use the same callback API, while an optimizer schedule is the appropriate interface when the rate must change after optimizer steps rather than after epochs.

Steps to set a Keras learning rate scheduler:

Create learning_rate_scheduler.py with the imports and deterministic regression data.

learning_rate_scheduler.py

import keras
import numpy as np
 
 
keras.utils.set_random_seed(21)
 
x_train = np.linspace(-1.0, 1.0, 48, dtype="float32").reshape(-1, 1)
y_train = (3.0 * x_train) - 0.25

Append the epoch-based scheduling function below the training data.
```
def schedule(epoch, learning_rate):
    if epoch < 2:
        return learning_rate
    return learning_rate * 0.5
```
The first callback invocation receives epoch index 0. Returning the current value for indexes 0 and 1 keeps the initial rate for the first two epochs.

Append the compiled regression model below schedule().

model = keras.Sequential(
    [
        keras.layers.Input(shape=(1,)),
        keras.layers.Dense(8, activation="relu"),
        keras.layers.Dense(1),
    ]
)
model.compile(
    optimizer=keras.optimizers.SGD(learning_rate=0.1),
    loss="mse",
)

The sample model, data, and loss stand in for the project's existing training objects; the optimizer's starting rate must remain aligned with the schedule.

Append the scheduler callback, training call, and final learning-rate assertion below the model.

lr_scheduler = keras.callbacks.LearningRateScheduler(schedule, verbose=1)
 
model.fit(
    x_train,
    y_train,
    epochs=5,
    batch_size=8,
    callbacks=[lr_scheduler],
    verbose=0,
)
 
final_learning_rate = float(
    keras.ops.convert_to_numpy(model.optimizer.learning_rate)
)
assert np.isclose(final_learning_rate, 0.0125)
print(f"final learning rate: {final_learning_rate:.4f}")

The project's callbacks list must contain lr_scheduler for every fit() call governed by the policy. The assertion fails if the five-epoch run does not leave the optimizer at 0.0125.

Run learning_rate_scheduler.py with the JAX backend for five scheduled epochs.

$ KERAS_BACKEND=jax python3 learning_rate_scheduler.py

Epoch 1: LearningRateScheduler setting learning rate to 0.10000000149011612.

Epoch 2: LearningRateScheduler setting learning rate to 0.10000000149011612.

Epoch 3: LearningRateScheduler setting learning rate to 0.05000000074505806.

Epoch 4: LearningRateScheduler setting learning rate to 0.02500000037252903.

Epoch 5: LearningRateScheduler setting learning rate to 0.012500000186264515.
final learning rate: 0.0125

The backend prefix can be KERAS_BACKEND=tensorflow or KERAS_BACKEND=torch when the installed Keras environment supports that backend. The five callback updates and final asserted value confirm that the scheduler controlled the optimizer across the complete run.

Related: How to set the Keras backend