TensorFlow 2
Adam Optimizer
Learning Rate Adjustment
Machine Learning
Neural Networks

How to change a learning rate for Adam in TF2?

Master System Design with Codemia

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

Introduction

In TensorFlow 2, changing Adam's learning rate can mean three different things: setting a fixed value when the optimizer is created, assigning a new value during training, or using a schedule that changes automatically over time. The best choice depends on whether the learning-rate change is manual, epoch-based, or step-based.

Set a Fixed Learning Rate at Optimizer Creation

The simplest case is a constant learning rate.

python
import tensorflow as tf

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

If you compile a Keras model with this optimizer, Adam will keep that value unless you change it later.

Change the Learning Rate Manually

If training is already running and you want to lower or raise the learning rate explicitly, update the optimizer's learning_rate variable.

python
1import tensorflow as tf
2
3optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
4
5optimizer.learning_rate.assign(1e-4)
6print(optimizer.learning_rate.numpy())

This is useful when you are driving training manually with a custom loop or when a callback decides the rate should change.

Use a Built-In Schedule

For most non-trivial training jobs, a schedule is cleaner than changing the value by hand. TensorFlow provides several learning-rate schedules that can be passed directly to Adam.

python
1import tensorflow as tf
2
3schedule = tf.keras.optimizers.schedules.ExponentialDecay(
4    initial_learning_rate=1e-3,
5    decay_steps=1000,
6    decay_rate=0.96,
7    staircase=True,
8)
9
10optimizer = tf.keras.optimizers.Adam(learning_rate=schedule)

Now the optimizer updates its effective learning rate automatically based on the current step count.

This is often preferable because the schedule becomes part of the optimizer definition instead of being scattered through training code.

Use a Callback in model.fit

If you train through model.fit, callbacks are another common way to adjust the learning rate.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(10, activation="relu"),
5    tf.keras.layers.Dense(1),
6])
7
8model.compile(
9    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
10    loss="mse",
11)
12
13callback = tf.keras.callbacks.ReduceLROnPlateau(
14    monitor="val_loss",
15    factor=0.5,
16    patience=3,
17)

ReduceLROnPlateau is especially useful when you want learning-rate changes to depend on validation behavior rather than a fixed step formula.

Custom Training Loops

In a custom training loop, you can combine manual assignment with any condition you like.

python
1for epoch in range(20):
2    if epoch == 10:
3        optimizer.learning_rate.assign(1e-4)
4
5    for x_batch, y_batch in dataset:
6        with tf.GradientTape() as tape:
7            predictions = model(x_batch, training=True)
8            loss = loss_fn(y_batch, predictions)
9
10        gradients = tape.gradient(loss, model.trainable_variables)
11        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

This gives full control when your training logic is more specialized than model.fit callbacks can express.

Common Pitfalls

The biggest pitfall is confusing the learning rate you set with the effective step behavior of Adam. Adam adapts parameter-wise updates internally, but the base learning rate still matters and still needs tuning.

Another issue is trying to assign() a new value when the learning rate was created as a non-assignable object or a schedule. If you passed a schedule into the optimizer, you should usually change the schedule design rather than overriding it in place.

Developers also sometimes expect ReduceLROnPlateau and a built-in schedule to cooperate automatically. In practice, you usually choose one mechanism per optimizer unless you are very deliberate about the interaction.

Finally, make sure you inspect the actual current value during debugging. A schedule-driven optimizer may not be using the initial number you remember setting earlier.

Summary

  • In TF2, Adam's learning rate can be fixed, assigned manually, or driven by a schedule.
  • Use learning_rate=... for a constant value, assign(...) for manual changes, and schedules for automatic step-based control.
  • In model.fit, callbacks such as ReduceLROnPlateau are often the easiest adaptive option.
  • In custom loops, manual assignment gives full control.
  • Be clear about whether the optimizer uses a plain variable or a schedule before trying to change it at runtime.

Course illustration
Course illustration

All Rights Reserved.