Python
NotImplementedError
Machine Learning
Learning Rate Schedule
Error Handling

NotImplementedError Learning rate schedule must override get_config

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

This error appears when a custom Keras or TensorFlow learning-rate schedule can be called during training but cannot be serialized. Keras expects schedule objects to describe how to rebuild themselves, and the minimum requirement for that is a proper get_config() method.

Why Keras Demands get_config()

Learning-rate schedules are often stored inside optimizers, and optimizers may be saved as part of the model or training configuration. For that to work, Keras needs a plain dictionary of constructor arguments it can use later to reconstruct the schedule.

If a custom schedule subclasses tf.keras.optimizers.schedules.LearningRateSchedule but does not implement get_config, Keras raises:

  • the schedule can run
  • the schedule cannot be serialized
  • Keras stops and asks you to define the missing config method

That is what the NotImplementedError is really about.

A Broken Custom Schedule

This custom schedule computes a learning rate, but it cannot be saved correctly:

python
1import tensorflow as tf
2
3
4class WarmupDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
5    def __init__(self, initial_lr, warmup_steps):
6        self.initial_lr = initial_lr
7        self.warmup_steps = warmup_steps
8
9    def __call__(self, step):
10        step = tf.cast(step, tf.float32)
11        warmup_steps = tf.cast(self.warmup_steps, tf.float32)
12        return self.initial_lr * tf.minimum(1.0, step / warmup_steps)

This may seem fine until Keras tries to serialize the optimizer or inspect the schedule configuration.

The Correct Fix

Add get_config() and return the arguments needed to recreate the object:

python
1import tensorflow as tf
2
3
4class WarmupDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
5    def __init__(self, initial_lr, warmup_steps):
6        self.initial_lr = initial_lr
7        self.warmup_steps = warmup_steps
8
9    def __call__(self, step):
10        step = tf.cast(step, tf.float32)
11        warmup_steps = tf.cast(self.warmup_steps, tf.float32)
12        return self.initial_lr * tf.minimum(1.0, step / warmup_steps)
13
14    def get_config(self):
15        return {
16            "initial_lr": self.initial_lr,
17            "warmup_steps": self.warmup_steps,
18        }

Now Keras knows how to serialize the schedule parameters. In simple cases, that is all you need.

Use the Schedule in an Optimizer

Once the schedule is defined correctly, you can plug it into an optimizer:

python
1schedule = WarmupDecay(initial_lr=1e-3, warmup_steps=1000)
2optimizer = tf.keras.optimizers.Adam(learning_rate=schedule)
3
4model = tf.keras.Sequential(
5    [
6        tf.keras.layers.Input(shape=(10,)),
7        tf.keras.layers.Dense(32, activation="relu"),
8        tf.keras.layers.Dense(1),
9    ]
10)
11
12model.compile(optimizer=optimizer, loss="mse")

If the schedule class participates in saved model workflows or configuration exports, the missing-get_config problem is now gone.

Optional: Register for Cleaner Serialization

If you want smoother loading in larger projects, registering the class can help:

python
1import tensorflow as tf
2
3
4@tf.keras.utils.register_keras_serializable()
5class WarmupDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
6    def __init__(self, initial_lr, warmup_steps):
7        self.initial_lr = initial_lr
8        self.warmup_steps = warmup_steps
9
10    def __call__(self, step):
11        step = tf.cast(step, tf.float32)
12        warmup_steps = tf.cast(self.warmup_steps, tf.float32)
13        return self.initial_lr * tf.minimum(1.0, step / warmup_steps)
14
15    def get_config(self):
16        return {
17            "initial_lr": self.initial_lr,
18            "warmup_steps": self.warmup_steps,
19        }

This is not what fixes the error by itself. The crucial part is still get_config(). Registration just makes deserialization cleaner across save and load boundaries.

What get_config() Should Return

Return the same data needed by __init__. That usually means:

  • plain Python numbers
  • strings
  • booleans
  • other serializable config values

Do not return tensors or runtime-only objects unless you also know how they will be reconstructed. A good rule is that get_config() should describe the schedule, not the temporary state of a particular training run.

Common Pitfalls

The most common mistake is implementing __call__ and assuming that is enough because the schedule “works” during training. It works for execution, but not for serialization.

Another pitfall is returning non-serializable values from get_config(). If the config dictionary contains tensors or opaque objects, loading can fail later for a different reason.

It is also easy to forget that the config should match the constructor arguments. If __init__ needs a value that get_config() does not return, reconstruction becomes incomplete.

Finally, do not confuse registration with configuration. @register_keras_serializable is helpful, but it does not replace get_config().

Summary

  • The error means Keras can execute your custom schedule but cannot serialize it.
  • Fix it by implementing get_config() on your LearningRateSchedule subclass.
  • Return the constructor arguments needed to rebuild the object.
  • Keep the config dictionary plain and serializable.
  • Optional registration improves loading ergonomics, but get_config() is the actual requirement.

Course illustration
Course illustration

All Rights Reserved.