How to run a custom training loop in TensorFlow

Built-in Keras training handles most model updates, but research code and unusual batch logic sometimes need a training step that the project owns directly. A custom TensorFlow loop gives that control while still using Keras layers, optimizers, losses, and metrics.

The loop feeds a batched tf.data.Dataset into train_step(), records the forward pass with tf.GradientTape, applies optimizer updates, and runs validation with training=False. After the eager version behaves correctly, tf.function can compile stable tensor-only step functions to reduce repeated Python overhead.

Manual training also moves metric state, validation timing, and regularization losses into project code. Reset metrics at the start of each epoch, add model.losses inside the tape block when the model uses layer regularizers, and configure mixed-precision optimizers before the first gradient update.

1. Save the runnable custom loop demo as training_custom_loop_demo.py.

   training_custom_loop_demo.py

   import os
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
   tf.get_logger().setLevel("ERROR")
   tf.keras.utils.set_random_seed(7)
    
   features = tf.random.stateless_normal((320, 8), seed=(7, 11))
   score = (
       features[:, 0] * 0.8
       + features[:, 1] * 0.6
       - features[:, 2] * 0.4
       + features[:, 3] * 0.2
   )
   labels = tf.cast(score > 0, tf.float32)[:, None]
    
   x_train = features[:256]
   y_train = labels[:256]
   x_val = features[256:]
   y_val = labels[256:]
    
   train_ds = (
       tf.data.Dataset.from_tensor_slices((x_train, y_train))
       .shuffle(256, seed=7, reshuffle_each_iteration=True)
       .batch(32)
   )
   val_ds = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(32)
    
   model = tf.keras.Sequential(
       [
           tf.keras.layers.Input(shape=(8,)),
           tf.keras.layers.Dense(16, activation="relu"),
           tf.keras.layers.Dense(8, activation="relu"),
           tf.keras.layers.Dense(1, activation="sigmoid"),
       ]
   )
    
   loss_fn = tf.keras.losses.BinaryCrossentropy()
   optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
    
   train_loss = tf.keras.metrics.Mean(name="train_loss")
   train_accuracy = tf.keras.metrics.BinaryAccuracy(name="train_accuracy")
   val_loss = tf.keras.metrics.Mean(name="val_loss")
   val_accuracy = tf.keras.metrics.BinaryAccuracy(name="val_accuracy")
    
    
   @tf.function
   def train_step(features, labels):
       with tf.GradientTape() as tape:
           predictions = model(features, training=True)
           loss = loss_fn(labels, predictions)
           if model.losses:
               loss += tf.add_n(model.losses)
       gradients = tape.gradient(loss, model.trainable_variables)
       optimizer.apply_gradients(zip(gradients, model.trainable_variables))
       train_loss.update_state(loss)
       train_accuracy.update_state(labels, predictions)
    
    
   @tf.function
   def val_step(features, labels):
       predictions = model(features, training=False)
       loss = loss_fn(labels, predictions)
       val_loss.update_state(loss)
       val_accuracy.update_state(labels, predictions)
    
    
   for epoch in range(1, 5):
       train_loss.reset_state()
       train_accuracy.reset_state()
       val_loss.reset_state()
       val_accuracy.reset_state()
    
       for features_batch, labels_batch in train_ds:
           train_step(features_batch, labels_batch)
    
       for features_batch, labels_batch in val_ds:
           val_step(features_batch, labels_batch)
    
       print(
           f"Epoch {epoch}: "
           f"train_loss={train_loss.result():.4f} "
           f"train_accuracy={train_accuracy.result():.4f} "
           f"val_loss={val_loss.result():.4f} "
           f"val_accuracy={val_accuracy.result():.4f}"
       )
    
   predictions = model(x_val[:4], training=False)[:, 0]
   formatted = [round(float(value), 4) for value in predictions]
   print(f"sample_predictions={formatted}")

   Related: How to create a virtual environment for TensorFlow
   Related: How to install TensorFlow with pip

2. Run the demo and confirm that validation accuracy improves by the final epoch.

   $ python3 training_custom_loop_demo.py
   Epoch 1: train_loss=0.6656 train_accuracy=0.5781 val_loss=0.6188 val_accuracy=0.7031
   Epoch 2: train_loss=0.5664 train_accuracy=0.7539 val_loss=0.4943 val_accuracy=0.7812
   Epoch 3: train_loss=0.4398 train_accuracy=0.8125 val_loss=0.3508 val_accuracy=0.8906
   Epoch 4: train_loss=0.2977 train_accuracy=0.9297 val_loss=0.2412 val_accuracy=0.9375
   sample_predictions=[0.0104, 0.6858, 0.0001, 0.0763]

   val_accuracy=0.9375 and the printed sample predictions show that the training loop updated model weights and that the validation loop ran after the training batches.

3. Keep the forward pass, regularization losses, gradients, and optimizer update inside the project train_step() function.

   @tf.function
   def train_step(features, labels):
       with tf.GradientTape() as tape:
           predictions = model(features, training=True)
           loss = loss_fn(labels, predictions)
           if model.losses:
               loss += tf.add_n(model.losses)
       gradients = tape.gradient(loss, model.trainable_variables)
       optimizer.apply_gradients(zip(gradients, model.trainable_variables))
       train_loss.update_state(loss)
       train_accuracy.update_state(labels, predictions)

   training=True keeps layers such as Dropout and BatchNormalization in training mode during the forward pass. Add a loss-scaling optimizer before this step when the project uses mixed precision.
   Related: How to enable mixed precision in Keras

4. Run validation in a separate step function without applying gradients.

   @tf.function
   def val_step(features, labels):
       predictions = model(features, training=False)
       loss = loss_fn(labels, predictions)
       val_loss.update_state(loss)
       val_accuracy.update_state(labels, predictions)

   training=False uses inference behavior for validation while leaving the optimizer and trainable variables untouched.

5. Reset metric state before each epoch iterates over the training and validation datasets.

   for epoch in range(1, epochs + 1):
       train_loss.reset_state()
       train_accuracy.reset_state()
       val_loss.reset_state()
       val_accuracy.reset_state()
    
       for features_batch, labels_batch in train_ds:
           train_step(features_batch, labels_batch)
    
       for features_batch, labels_batch in val_ds:
           val_step(features_batch, labels_batch)

   Skipping reset_state() makes metrics accumulate across earlier epochs and hides whether the current epoch is improving.
   Related: How to resume Keras training from a checkpoint

6. Remove the temporary demo after the project loop reports training and validation metrics.

   $ rm training_custom_loop_demo.py