Custom Training Loops with GradientTape

When fit() isn't enough

Sometimes model.fit() isn't flexible enough — GANs need interleaved generator/discriminator updates, multi-task losses with custom weighting, gradient surgery for research. Custom training loops with GradientTape give you full control.

Always decorate your train_step with @tf.function outside the loop. The first call traces and compiles the graph; subsequent calls run the compiled version directly. This is what gives custom loops parity with fit()'s speed.

Keras metrics accumulate. Calling metric.result() returns the running average since the last reset_state(). If you forget to reset between epochs, your epoch-2 metrics include epoch-1 data and you'll wonder why numbers look weird. Reset after every epoch summary.

Code

Custom training loop — 정석 패턴·python

import tensorflow as tf
from tensorflow import keras

model = keras.Sequential([keras.Input(shape=(784,)),
                          keras.layers.Dense(256, activation='relu'),
                          keras.layers.Dense(10)])
optimizer = keras.optimizers.Adam(1e-3)
loss_fn   = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

train_loss = keras.metrics.Mean(name='train_loss')
train_acc  = keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
val_loss   = keras.metrics.Mean(name='val_loss')
val_acc    = keras.metrics.SparseCategoricalAccuracy(name='val_accuracy')

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss = loss_fn(y, logits)
        loss += sum(model.losses)         # add regularization losses
    grads = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    train_loss.update_state(loss)
    train_acc.update_state(y, logits)
    return loss

@tf.function
def val_step(x, y):
    logits = model(x, training=False)
    val_loss.update_state(loss_fn(y, logits))
    val_acc.update_state(y, logits)

EPOCHS = 10
for epoch in range(EPOCHS):
    for x_b, y_b in train_ds:
        train_step(x_b, y_b)
    for x_b, y_b in val_ds:
        val_step(x_b, y_b)

    print(f"Epoch {epoch+1}/{EPOCHS} — "
          f"loss: {train_loss.result():.4f}  acc: {train_acc.result():.4f}  "
          f"val_loss: {val_loss.result():.4f}  val_acc: {val_acc.result():.4f}")

    # IMPORTANT: reset between epochs
    train_loss.reset_state(); train_acc.reset_state()
    val_loss.reset_state();   val_acc.reset_state()

Custom Training Loops with GradientTape

When fit() isn't enough

Code

Progress

Comments 0