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.

1. 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

2. 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.

3. 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.

4. 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.

   Related: How to compile a model in Keras

5. 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.

   Related: How to create a custom callback in Keras

6. Remove the smoke test file after copying the scheduler into the project code.

   $ rm learning_rate_scheduler_demo.py