Introduction

Accessing model predictions at the end of each epoch is useful for debugging learning dynamics, detecting overfitting early, and building custom monitoring workflows beyond standard metrics. In Keras, the cleanest way to do this is a callback that runs model.predict in on_epoch_end.

However, naïvely predicting on a full validation set every epoch can be expensive and distort training time. A production-ready approach samples data, stores outputs efficiently, and avoids memory growth from accumulating large arrays.

Core Sections

1. Build a callback to collect per-epoch predictions

1import tensorflow as tf
2import numpy as np
3
4class EpochPredictionCallback(tf.keras.callbacks.Callback):
5    def __init__(self, x_eval, y_eval=None, max_samples=256):
6        super().__init__()
7        self.x_eval = x_eval[:max_samples]
8        self.y_eval = None if y_eval is None else y_eval[:max_samples]
9        self.history_preds = []
10
11    def on_epoch_end(self, epoch, logs=None):
12        preds = self.model.predict(self.x_eval, verbose=0)
13        self.history_preds.append(preds)
14        print(f"Epoch {epoch+1}: stored predictions shape={preds.shape}")

Use with training:

callback = EpochPredictionCallback(x_val, y_val, max_samples=128)
model.fit(x_train, y_train, epochs=10, callbacks=[callback], validation_data=(x_val, y_val))

Now callback.history_preds[i] contains predictions from epoch i+1.

2. Store compact summaries instead of full arrays

For long runs, storing full predictions each epoch may consume large memory. Store summaries:

1class EpochStatsCallback(tf.keras.callbacks.Callback):
2    def __init__(self, x_eval):
3        self.x_eval = x_eval
4        self.stats = []
5
6    def on_epoch_end(self, epoch, logs=None):
7        p = self.model.predict(self.x_eval, verbose=0)
8        self.stats.append({
9            "epoch": epoch + 1,
10            "mean": float(np.mean(p)),
11            "std": float(np.std(p)),
12            "min": float(np.min(p)),
13            "max": float(np.max(p)),
14        })

This drastically reduces memory footprint while still tracking dynamics.

3. Save per-epoch predictions to disk

If you need full outputs for later analysis, save and release memory immediately.

1from pathlib import Path
2
3class SavePredictionsCallback(tf.keras.callbacks.Callback):
4    def __init__(self, x_eval, out_dir="epoch_preds"):
5        self.x_eval = x_eval
6        self.out_dir = Path(out_dir)
7        self.out_dir.mkdir(parents=True, exist_ok=True)
8
9    def on_epoch_end(self, epoch, logs=None):
10        preds = self.model.predict(self.x_eval, verbose=0)
11        np.save(self.out_dir / f"pred_epoch_{epoch+1:03d}.npy", preds)

This is better for large validation sets or long epoch counts.

4. Example visualization of prediction drift

1import matplotlib.pyplot as plt
2
3first_sample = [preds[0, 0] for preds in callback.history_preds]
4plt.plot(range(1, len(first_sample)+1), first_sample)
5plt.xlabel("Epoch")
6plt.ylabel("Prediction for sample[0]")
7plt.title("Prediction evolution")
8plt.show()

This can reveal instability even when aggregate validation loss looks acceptable.

Common Pitfalls

• Predicting on the entire validation dataset every epoch and significantly slowing training.
• Appending large prediction arrays in memory without bounds, causing RAM growth.
• Forgetting verbose=0 in callbacks, which clutters logs and slows pipelines.
• Comparing predictions across epochs without keeping evaluation subset fixed.
• Running expensive post-processing inside callback on training thread, increasing epoch time variance.

Summary

In Keras, per-epoch predictions are best captured with custom callbacks using on_epoch_end. Keep the evaluation subset stable, control memory by summarizing or writing to disk, and separate analysis from training critical path. With these patterns, you can inspect model behavior deeply without destabilizing training performance.

For classification tasks, raw predictions are often less informative than derived diagnostics. At each epoch you can compute confusion matrices, calibration curves, or top-k changes on a fixed holdout subset. This reveals whether predictions are becoming more confident, merely shifting thresholds, or oscillating due to unstable training. Logging those diagnostics to TensorBoard or experiment trackers provides richer insight than scalar loss/accuracy alone.

When training on distributed setups, ensure callbacks run in a way compatible with strategy scope and data sharding. Some callback patterns that work in single-process notebooks can become expensive or inconsistent under distributed training. A good pattern is to keep epoch-end prediction subsets small and deterministic, then aggregate summaries instead of full tensors. This keeps monitoring informative without significantly increasing epoch duration.

When storing prediction snapshots, include epoch number, model hash, and dataset slice identifier so later analysis remains traceable and reproducible.

This keeps monitoring useful while preserving training throughput.

Use this pattern to balance insight, reproducibility, and cost.