Custom Training Loop

Code

import tensorflow as tf
import numpy as np

model = tf.keras.Sequential([tf.keras.layers.Dense(2, input_shape=(4,))])
optimizer = tf.keras.optimizers.Adam(1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

train_ds = tf.data.Dataset.from_tensor_slices(
    (np.random.rand(200, 4).astype("float32"), np.random.randint(0, 2, (200,)))
).batch(32)

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss = loss_fn(y, logits)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

for epoch in range(5):
    total = 0.0
    for x, y in train_ds:
        total += float(train_step(x, y))
    print(f"epoch {epoch} loss={total:.4f}")

Explanation

A custom loop gives full control: GradientTape computes gradients of the loss with respect to trainable variables, and apply_gradients updates them. Wrapping the step in @tf.function graphs it for major speedups over eager mode. Custom loops are essential when the standard fit() cannot express the logic.

More TensorFlow Snippets