Custom Models and Training Loops (TensorFlow)

model.fit() handles batching, shuffling, metrics, callbacks and validation. When you need something it does not support — a second optimiser, an adversarial inner loop, a loss that reads intermediate activations — you write the loop yourself.

That trade is usually described as “flexibility for verbosity”. The real cost is different, and it is measurable:

Run	Final accuracy	Seconds
`model.fit`	0.9100	3.0
Custom loop, `@tf.function`	0.9110	2.3
Custom loop, eager	0.9110	12.5

Same data, same seed, same batches, no shuffling. The hand-written loop is correct — and the only difference between the fast version and the slow one is a single decorator.

What you’ll learn

How to check a hand-written loop is equivalent to fit before trusting anything it produces.
Why tf.function is worth 9.9× per step, and what it actually changes.
Why the custom loop beat fit on wall-clock, and why that is not a reason to abandon fit.
How to write a loss and a metric by hand, and the averaging bug that catches people.

The loop

Everything fit() does, unrolled

optimizer = keras.optimizers.Adam(1e-3)
loss_function = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
 
@tf.function                                   # <- the 9.9x
def step(images, labels):
    with tf.GradientTape() as tape:
        logits = model(images, training=True)  # training=True matters for
        loss = loss_function(labels, logits)   # dropout and batch norm
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss
 
for epoch in range(EPOCHS):
    for begin in range(0, len(images) - BATCH + 1, BATCH):
        step(images[begin:begin + BATCH], labels[begin:begin + BATCH])

Three details in that block are easy to get wrong and hard to notice:

training=True switches dropout on and batch normalisation into batch-statistics mode. Omit it and the model trains in inference mode, which usually still converges — slightly worse, with no error.
The tape only records what happens inside it. A loss computed outside the with block has no gradient path, and tape.gradient returns None rather than raising.
apply_gradients takes pairs, so a mismatch between the gradient list and the variable list fails silently in the sense that it applies something, just not what you meant.

Check the loop before trusting it

Left: test accuracy per epoch for model.fit and the custom loop, two nearly overlapping curves rising from about 0.82 to 0.91. Right: bars of the absolute difference between them at each epoch, all at or below 0.0045. — The comparison only means something because shuffling is off on both sides — with shuffling the two would see different batches and any difference would be unattributable. The residual gap of 0.0010 to 0.0045 comes from floating-point ordering inside the two implementations, not from the loop being wrong. Establishing this first is what makes the timing comparison that follows interpretable.

Epoch	`fit`	Custom	Difference
1	0.8230	0.8245	0.0015
2	0.8765	0.8810	0.0045
3	0.8915	0.8930	0.0015
6	0.9100	0.9110	0.0010

This step is not ceremony. A custom loop that quietly skips the last partial batch, applies gradients twice, or forgets training=True will still produce a curve that goes up — and you will spend the next week attributing your research result to the wrong cause.

What `tf.function` does

Left: two bars of milliseconds per training step, 31.21 for eager and 3.14 with tf.function. Right: horizontal bars of total seconds for six epochs — model.fit at 3.0, the compiled custom loop at 2.3, and the eager custom loop at 12.5. — 31.21 ms against 3.14 ms per step. Eager mode dispatches every operation from Python one at a time; tf.function traces the whole step once into a graph and then executes it without returning to the interpreter between operations. On a small model the per-operation Python overhead dominates completely, which is why the ratio is so large here and would shrink on a model whose individual operations are expensive.

Tracing has consequences worth knowing before you rely on it:

Python side effects run only during tracing. A print() inside a tf.function fires once, not every step. Use tf.print if you need it every time.
A new input shape triggers a re-trace. A final partial batch of a different size compiles a second graph; reduce_retracing=True relaxes this.
Python conditionals are baked in. if training: on a Python boolean is decided at trace time; on a tensor it needs tf.cond.

Why the custom loop beat `fit`

2.3 seconds against 3.0. That is real but it is not a reason to stop using fit — the gap is everything fit does that this loop does not: computing running metrics each batch, evaluating on the validation set every epoch, and running the callback machinery.

The honest framing is that fit costs about 0.7 seconds of features over six epochs here, and those features are ones you would otherwise write yourself — usually less carefully. Write the loop when you need control, not for speed.

diagram Diagram mermaid

  flowchart TD
  F["model.fit"] --> A["batching, shuffling, metrics,
callbacks, validation"]
  A --> R1["3.0s, 0.9100"]
  C["custom loop"] --> B{"tf.function?"}
  B -->|"yes"| G["traced once, runs as a graph"]
  B -->|"no"| E["every op dispatched from Python"]
  G --> R2["2.3s, 0.9110"]
  E --> R3["12.5s, 0.9110"]
  R3 -.->|"9.9x per step"| W["the same maths, 5.4x the wall-clock"]

Custom losses and metrics

A hand-written loss should reproduce the built-in it replaces before you start changing it:

Left: paired bars comparing built-in and hand-written loss and accuracy on the same predictions, matching to six decimal places. Right: per-batch accuracy for four batches of sizes 512, 512, 512 and 464, with a dashed line for the correct overall accuracy and a dotted line for the mean of the batch means, which sits slightly higher. — The loss matched exactly (difference 0.00e+00) and the accuracy to 1.34e-08, which is float32 summation order. The right panel is the bug worth knowing: averaging four per-batch accuracies gives 0.911116 against the correct 0.911000, because the final batch of 464 is weighted as heavily as the three batches of 512.

The averaging bug

# Wrong: every batch counts equally, whatever its size.
accuracy = np.mean([batch_accuracy(b) for b in batches])
 
# Right: accumulate counts, divide once.
metric = keras.metrics.SparseCategoricalAccuracy()
for images, labels in batches:
    metric.update_state(labels, model(images, training=False))
accuracy = metric.result()

The error here is 0.000116 — small, and always in the direction of whichever batches happen to be short. It is exactly the kind of discrepancy that gets attributed to “evaluation noise” when two implementations disagree.

Keras metric objects exist to avoid this: they accumulate a numerator and a denominator across update_state calls and divide once at result().

sketch Eager against traced p5.js

Drag the operation count. Each block is one operation; eager pays Python dispatch for every one, while a traced graph pays it once.

sketch The measured table, ranked p5.js

Click a column to rank every row by it. The bars are that column's values and the highest and lowest are computed from the numbers, not written in.

Pitfalls

Forgetting training=True in the forward pass. Dropout and batch norm silently run in inference mode and the model still trains, slightly worse.
Computing the loss outside the GradientTape block. tape.gradient returns None instead of raising.
Leaving tf.function off. 31.21 ms against 3.14 ms per step here — 9.9×.
Expecting print to fire every step inside a traced function. It runs at trace time only; tf.print runs every step.
Comparing a custom loop with fit while shuffling is on. The two see different batches and the comparison means nothing.
Averaging per-batch metrics. 0.911116 against the correct 0.911000, because the short final batch is over-weighted.
Assuming the speed advantage is the reason to write a loop. It came from skipping metrics, validation and callbacks — features, not overhead.

Recap

A custom loop reproduced fit to within 0.0010 accuracy with the same seed, batches and no shuffling — establish that before concluding anything else.
tf.function traces the step into a graph: 31.21 ms → 3.14 ms per step, 12.5s → 2.3s over six epochs.
Tracing means Python side effects run once, new shapes re-trace, and Python conditionals are frozen at trace time.
The custom loop’s 2.3s against fit’s 3.0s is the cost of the features fit provides, not waste.
A hand-written loss matched the built-in exactly; accuracy matched to 1.34e-08.
Averaging per-batch accuracies is wrong whenever the last batch is short — 0.911116 against 0.911000.

One loop on one device is the simple case. Spreading the same computation across several devices introduces a different set of problems: Distributed Training with tf.distribute.

pch.quizTag pch.quizDefaultTitle

The custom loop took 12.5s eager and 2.3s with @tf.function, for identical results. What does the decorator change?

pch.quizShowAnswer

B — It traces the whole step into a graph once, so operations execute without returning to the Python interpreter between each one — eager mode dispatches every operation individually — Both runs reached 0.9110, so the arithmetic is identical. On a small model the per-operation Python overhead dominates, which is why the ratio is 9.9x here.
Why must shuffling be disabled when comparing a custom loop against model.fit?

pch.quizShowAnswer

B — Otherwise the two runs see different batches in different orders, so any difference in accuracy cannot be attributed to the loop implementation — With shuffling off, the remaining 0.0010 to 0.0045 gap is attributable to floating-point summation order rather than to a bug in the loop.
Averaging four per-batch accuracies gave 0.911116 while the correct value was 0.911000. What causes the gap?

pch.quizShowAnswer

B — The batches were sizes 512, 512, 512 and 464, and averaging the four means weights the short final batch as heavily as the full ones — Keras metric objects avoid this by accumulating a numerator and denominator across update_state calls and dividing once at result().
What happens if you forget training=True in the forward pass of a custom loop?

pch.quizShowAnswer

B — Dropout is disabled and batch normalisation uses its moving statistics instead of batch statistics — the model still trains and usually converges slightly worse, with no warning — It is a silent failure, which is why establishing equivalence with fit before trusting a loop is worth the effort.
The custom loop finished in 2.3s against fit's 3.0s. Is that a reason to prefer custom loops generally?

pch.quizShowAnswer

B — No — the difference is the work fit does that the loop skips: per-batch metrics, per-epoch validation, and the callback machinery. Write a loop for control, not for speed — Those are features you would otherwise implement yourself, usually less carefully — as the per-batch averaging bug on this page shows.

🧪 Try It Yourself

pch.coffeeTagline

pch.coffeeCta

pch.feedbackHeading

pch.feedbackSubheading