How to run a custom training loop in TensorFlow

Built-in Keras training covers common model updates, but some algorithms need direct control over each batch, gradient calculation, validation pass, or metric boundary. A custom TensorFlow loop exposes those moments while retaining Keras models, losses, optimizers, and metrics.

Each training batch passes through the model with training=True inside tf.GradientTape. The tape connects the computed loss to the trainable variables, while the optimizer applies the resulting gradients; layer regularization terms from model.losses remain part of that loss.

Validation uses training=False and never applies gradients. Stateful metrics need reset_state() at each epoch boundary, and a useful smoke test should fail when validation loss does not improve or validation accuracy stays below the chosen threshold.

Steps to run a custom training loop in TensorFlow:

Create training_custom_loop.py with reproducible training and validation datasets.

training_custom_loop.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((512, 4), seed=(7, 11))
scores = 1.5 * features[:, 0] - features[:, 1] + 0.5 * features[:, 2]
labels = tf.cast(scores > 0, tf.float32)[:, None]
 
train_dataset = (
    tf.data.Dataset.from_tensor_slices((features[:448], labels[:448]))
    .shuffle(448, seed=7, reshuffle_each_iteration=True)
    .batch(32)
)
validation_dataset = tf.data.Dataset.from_tensor_slices(
    (features[448:], labels[448:])
).batch(32)

The fixed random seeds make the small classification problem repeatable. Project datasets can use the same batched features, labels structure.
Related: How to build a TensorFlow dataset input pipeline
Related: How to set a random seed in TensorFlow

Append the model configuration and metric state to training_custom_loop.py.

model = tf.keras.Sequential(
    [
        tf.keras.layers.Input(shape=(4,)),
        tf.keras.layers.Dense(
            8,
            activation="relu",
            kernel_regularizer=tf.keras.regularizers.L2(1e-4),
        ),
        tf.keras.layers.Dense(1, activation="sigmoid"),
    ]
)
loss_function = tf.keras.losses.BinaryCrossentropy()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.02)
 
train_loss = tf.keras.metrics.Mean(name="train_loss")
train_accuracy = tf.keras.metrics.BinaryAccuracy(name="train_accuracy")
validation_loss = tf.keras.metrics.Mean(name="validation_loss")
validation_accuracy = tf.keras.metrics.BinaryAccuracy(name="validation_accuracy")

The Mean metrics accumulate batch losses, while BinaryAccuracy accumulates predictions across the current epoch.

Append the compiled gradient-update function to training_custom_loop.py.

@tf.function
def train_step(batch_features, batch_labels):
    with tf.GradientTape() as tape:
        predictions = model(batch_features, training=True)
        loss = loss_function(batch_labels, predictions)
        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(batch_labels, predictions)

The model contains an L2 regularizer, so tf.add_n(model.losses) adds that penalty before gradients are calculated. Custom loops using mixed_float16 need loss scaling through LossScaleOptimizer, while mixed_bfloat16 does not need loss scaling.
Related: How to compute gradients with TensorFlow GradientTape
Related: How to enable mixed precision in Keras
Related: How to compile a function with tf.function in TensorFlow

Append the compiled validation function to training_custom_loop.py.

@tf.function
def validation_step(batch_features, batch_labels):
    predictions = model(batch_features, training=False)
    loss = loss_function(batch_labels, predictions)
    loss += tf.add_n(model.losses)
    validation_loss.update_state(loss)
    validation_accuracy.update_state(batch_labels, predictions)

training=False selects inference behavior for layers such as Dropout and BatchNormalization without creating an optimizer update.

Complete training_custom_loop.py with the epoch loop and fail-capable result checks.

initial_validation_loss = None
 
for epoch in range(1, 6):
    train_loss.reset_state()
    train_accuracy.reset_state()
    validation_loss.reset_state()
    validation_accuracy.reset_state()
 
    for batch_features, batch_labels in train_dataset:
        train_step(batch_features, batch_labels)
 
    for batch_features, batch_labels in validation_dataset:
        validation_step(batch_features, batch_labels)
 
    if initial_validation_loss is None:
        initial_validation_loss = tf.identity(validation_loss.result())
 
    print(
        f"epoch={epoch} "
        f"train_loss={train_loss.result():.4f} "
        f"train_accuracy={train_accuracy.result():.4f} "
        f"validation_loss={validation_loss.result():.4f} "
        f"validation_accuracy={validation_accuracy.result():.4f}"
    )
 
tf.debugging.assert_less(validation_loss.result(), initial_validation_loss)
tf.debugging.assert_greater(validation_accuracy.result(), 0.90)
 
print(
    "custom_loop_check=passed "
    f"optimizer_steps={int(optimizer.iterations)} "
    f"validation_accuracy={validation_accuracy.result():.4f}"
)

Skipping reset_state() makes each metric include earlier epochs, which hides the current epoch's loss and accuracy.

Compare the completed training_custom_loop.py file with the consolidated version.

training_custom_loop.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((512, 4), seed=(7, 11))
scores = 1.5 * features[:, 0] - features[:, 1] + 0.5 * features[:, 2]
labels = tf.cast(scores > 0, tf.float32)[:, None]
 
train_dataset = (
    tf.data.Dataset.from_tensor_slices((features[:448], labels[:448]))
    .shuffle(448, seed=7, reshuffle_each_iteration=True)
    .batch(32)
)
validation_dataset = tf.data.Dataset.from_tensor_slices(
    (features[448:], labels[448:])
).batch(32)
 
model = tf.keras.Sequential(
    [
        tf.keras.layers.Input(shape=(4,)),
        tf.keras.layers.Dense(
            8,
            activation="relu",
            kernel_regularizer=tf.keras.regularizers.L2(1e-4),
        ),
        tf.keras.layers.Dense(1, activation="sigmoid"),
    ]
)
loss_function = tf.keras.losses.BinaryCrossentropy()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.02)
 
train_loss = tf.keras.metrics.Mean(name="train_loss")
train_accuracy = tf.keras.metrics.BinaryAccuracy(name="train_accuracy")
validation_loss = tf.keras.metrics.Mean(name="validation_loss")
validation_accuracy = tf.keras.metrics.BinaryAccuracy(name="validation_accuracy")
 
 
@tf.function
def train_step(batch_features, batch_labels):
    with tf.GradientTape() as tape:
        predictions = model(batch_features, training=True)
        loss = loss_function(batch_labels, predictions)
        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(batch_labels, predictions)
 
 
@tf.function
def validation_step(batch_features, batch_labels):
    predictions = model(batch_features, training=False)
    loss = loss_function(batch_labels, predictions)
    loss += tf.add_n(model.losses)
    validation_loss.update_state(loss)
    validation_accuracy.update_state(batch_labels, predictions)
 
 
initial_validation_loss = None
 
for epoch in range(1, 6):
    train_loss.reset_state()
    train_accuracy.reset_state()
    validation_loss.reset_state()
    validation_accuracy.reset_state()
 
    for batch_features, batch_labels in train_dataset:
        train_step(batch_features, batch_labels)
 
    for batch_features, batch_labels in validation_dataset:
        validation_step(batch_features, batch_labels)
 
    if initial_validation_loss is None:
        initial_validation_loss = tf.identity(validation_loss.result())
 
    print(
        f"epoch={epoch} "
        f"train_loss={train_loss.result():.4f} "
        f"train_accuracy={train_accuracy.result():.4f} "
        f"validation_loss={validation_loss.result():.4f} "
        f"validation_accuracy={validation_accuracy.result():.4f}"
    )
 
tf.debugging.assert_less(validation_loss.result(), initial_validation_loss)
tf.debugging.assert_greater(validation_accuracy.result(), 0.90)
 
print(
    "custom_loop_check=passed "
    f"optimizer_steps={int(optimizer.iterations)} "
    f"validation_accuracy={validation_accuracy.result():.4f}"
)

Run training_custom_loop.py to confirm that the loop updates weights and passes the validation checks.

$ python3 training_custom_loop.py
epoch=1 train_loss=0.7179 train_accuracy=0.4911 validation_loss=0.5992 validation_accuracy=0.7500
epoch=2 train_loss=0.5246 train_accuracy=0.8371 validation_loss=0.4398 validation_accuracy=0.8750
epoch=3 train_loss=0.3377 train_accuracy=0.9241 validation_loss=0.2689 validation_accuracy=0.9219
epoch=4 train_loss=0.1900 train_accuracy=0.9821 validation_loss=0.1814 validation_accuracy=0.9531
epoch=5 train_loss=0.1256 train_accuracy=0.9866 validation_loss=0.1355 validation_accuracy=0.9375
custom_loop_check=passed optimizer_steps=70 validation_accuracy=0.9375

The final line appears only after validation loss falls below its first-epoch value and validation accuracy exceeds 0.90.

Author: Mohd Shakir Zakaria
Mohd Shakir Zakaria is a cloud architect with deep roots in software development and open-source advocacy. Certified in AWS, Red Hat, VMware, ITIL, and Linux, he specializes in designing and managing robust cloud and on-premises infrastructures.