How to set a learning rate scheduler in Keras

Learning-rate changes shape how quickly a Keras model updates its weights during training. A scheduler is useful when a run should start with larger optimizer steps and shrink them after warm-up epochs instead of restarting training with a manually edited optimizer setting.

LearningRateScheduler is an epoch-level callback. It receives the epoch index and the current optimizer learning rate at the start of each epoch, returns the next value, and lets model.fit() continue with that rate for the epoch that is about to run.

The smoke test uses a tiny regression model, SGD with an initial learning rate of 0.1, and a schedule that keeps the first two epochs unchanged before halving the value each epoch. The printed rates are the proof of the callback behavior; the final loss can vary slightly by backend and hardware.

Steps to set a Keras learning rate scheduler:

Open a terminal in the Python environment where Keras imports with the backend used by the training project.

Set the Keras backend before Python imports keras when the project uses JAX or PyTorch instead of the default backend.
Related: How to set the Keras backend

Save the scheduler smoke test as

learning_rate_scheduler_demo.py

learning_rate_scheduler_demo.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
 
 
class LearningRatePrinter(keras.callbacks.Callback):
    def on_train_begin(self, logs=None):
        self.rates = []
 
    def on_epoch_begin(self, epoch, logs=None):
        lr = float(keras.ops.convert_to_numpy(self.model.optimizer.learning_rate))
        self.rates.append(lr)
        print(f"epoch={epoch + 1} lr={lr:.4f}")
 
 
def schedule(epoch, lr):
    if epoch < 2:
        return lr
    return lr * 0.5
 
 
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",
)
 
lr_printer = LearningRatePrinter()
lr_scheduler = keras.callbacks.LearningRateScheduler(schedule)
 
history = model.fit(
    x_train,
    y_train,
    epochs=5,
    batch_size=8,
    callbacks=[lr_scheduler, lr_printer],
    verbose=0,
)
 
rates = ", ".join(f"{rate:.4f}" for rate in lr_printer.rates)
print(f"rates seen: [{rates}]")
print(f"final loss: {history.history['loss'][-1]:.4f}")

The schedule() function receives a zero-based epoch index and the current learning rate. Return the original value during warm-up epochs, then return the adjusted value for later epochs.

Run the smoke test.
```
$ KERAS_BACKEND=jax python3 learning_rate_scheduler_demo.py
epoch=1 lr=0.1000
epoch=2 lr=0.1000
epoch=3 lr=0.0500
epoch=4 lr=0.0250
epoch=5 lr=0.0125
rates seen: [0.1000, 0.1000, 0.0500, 0.0250, 0.0125]
final loss: 0.0325
```
Use the KERAS_BACKEND prefix only when the smoke-test environment needs an explicit backend. A run that prints unchanged rates for epochs 1 and 2, then halved rates for epochs 3 through 5 confirms that the scheduler is being applied.

Copy the scheduler callback into the project training call.

def schedule(epoch, lr):
    if epoch < 2:
        return lr
    return lr * 0.5
 
 
lr_scheduler = keras.callbacks.LearningRateScheduler(schedule)
 
history = model.fit(
    train_features,
    train_labels,
    validation_data=(val_features, val_labels),
    epochs=50,
    callbacks=[lr_scheduler],
)

Use an optimizer schedule such as keras.optimizers.schedules.ExponentialDecay instead when the learning rate should change by optimizer step rather than by epoch.

Add a temporary printer callback to the next project run when the training logs do not already expose the learning rate.

class LearningRatePrinter(keras.callbacks.Callback):
    def on_epoch_begin(self, epoch, logs=None):
        lr = float(keras.ops.convert_to_numpy(self.model.optimizer.learning_rate))
        print(f"epoch={epoch + 1} lr={lr:.6f}")
 
 
history = model.fit(
    train_features,
    train_labels,
    validation_data=(val_features, val_labels),
    epochs=50,
    callbacks=[lr_scheduler, LearningRatePrinter()],
)

Keep the scheduler before callbacks that read the learning rate at epoch start so those callbacks see the updated value.

Remove the smoke test file after copying the scheduler into the project code.
```
$ rm learning_rate_scheduler_demo.py
```