TopolearnCustom loops are needed for GANs, reinforcement learning, unusual losses and research where fit() does not fit.
model.fit() covers standard training. When you need control over each step, write the loop yourself from the pieces you already know:
1. Forward pass inside a tf.GradientTape, calling the model with training=True so dropout and batch norm behave correctly. 2. Compute the loss. 3. grads = tape.gradient(loss, model.trainable_variables). 4. optimizer.apply_gradients(zip(grads, model.trainable_variables)).
Wrap the step in @tf.function. Eager execution is convenient for debugging but much slower for many small ops; the compiled step usually runs several times faster. Keep the Python epoch loop outside the function and iterate a tf.data.Dataset so every batch has the same signature (use drop_remainder=True if a smaller last batch would cause a retrace).
Track progress with Keras metrics (tf.keras.metrics.Mean, SparseCategoricalAccuracy), calling update_state per step and result() / reset_state() per epoch.
A middle ground also exists: subclass keras.Model and override train_step, keeping fit()'s callbacks and progress bars while customising the step. That is often the best option when only the step logic differs.
Example
import tensorflow as tf
model = tf.keras.Sequential([tf.keras.Input((4,)),
tf.keras.layers.Dense(16, activation="relu"),
tf.keras.layers.Dense(1)])
optimizer = tf.keras.optimizers.Adam(1e-2)
loss_fn = tf.keras.losses.MeanSquaredError()
x = tf.random.normal([256, 4])
y = tf.reduce_sum(x, axis=1, keepdims=True)
ds = tf.data.Dataset.from_tensor_slices((x, y)).shuffle(256).batch(32)
@tf.function
def train_step(xb, yb):
with tf.GradientTape() as tape:
pred = model(xb, training=True)
loss = loss_fn(yb, pred)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
for epoch in range(30):
for xb, yb in ds:
loss = train_step(xb, yb)
print(f"last batch loss {loss.numpy():.3f}") # small, well under 0.1Remember
- Tape, loss,
tape.gradient,apply_gradients - Call the model with
training=Trueinside the step - Compile the step with
@tf.function, keep the epoch loop in Python
Watch out for
- Forgetting
training=True, so dropout and batch norm run in inference mode during training.