Introduction

Keras callbacks are great for logging, checkpointing, and early stopping, but they do not automatically hand you the full input batch during training. If you need access to batch tensors in TensorFlow 2, the clean solution is usually to expose the specific values you need from train_step and read them inside the callback.

Why callbacks do not receive x and y directly

Methods such as on_train_batch_begin and on_train_batch_end receive the batch index and a logs dictionary. That dictionary contains metrics and loss values, not the raw tensors that were fed into the model.

This design keeps callbacks generic, but it means code like "print the mean of the current input batch" cannot be done from a callback alone. You need a bridge between the training step and the callback.

A practical pattern: custom train_step

The standard pattern is:

1. subclass tf.keras.Model
2. override train_step
3. compute the batch-level values you care about
4. return them in logs
5. read them in the callback

Here is a minimal runnable example:

1import tensorflow as tf
2import numpy as np
3
4class InspectableModel(tf.keras.Model):
5    def train_step(self, data):
6        x, y = data
7
8        with tf.GradientTape() as tape:
9            y_pred = self(x, training=True)
10            loss = self.compiled_loss(y, y_pred)
11
12        gradients = tape.gradient(loss, self.trainable_variables)
13        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
14        self.compiled_metrics.update_state(y, y_pred)
15
16        logs = {metric.name: metric.result() for metric in self.metrics}
17        logs["loss"] = loss
18        logs["input_mean"] = tf.reduce_mean(x)
19        logs["target_mean"] = tf.reduce_mean(y)
20        return logs
21
22class BatchLogger(tf.keras.callbacks.Callback):
23    def on_train_batch_end(self, batch, logs=None):
24        logs = logs or {}
25        print(
26            f"batch={batch} "
27            f"loss={float(logs['loss']):.4f} "
28            f"input_mean={float(logs['input_mean']):.4f} "
29            f"target_mean={float(logs['target_mean']):.4f}"
30        )
31
32x = np.random.randn(64, 4).astype("float32")
33y = (x.sum(axis=1, keepdims=True) > 0).astype("float32")
34
35inputs = tf.keras.Input(shape=(4,))
36outputs = tf.keras.layers.Dense(1, activation="sigmoid")(inputs)
37model = InspectableModel(inputs, outputs)
38
39model.compile(
40    optimizer="adam",
41    loss="binary_crossentropy",
42    metrics=["accuracy"],
43)
44
45model.fit(x, y, epochs=1, batch_size=16, callbacks=[BatchLogger()], verbose=0)

The callback still does not receive raw batch tensors directly, but it receives exactly the batch-derived information you chose to publish.

If you really need the full batch tensors

Sometimes summaries are not enough. You may want to inspect misclassified samples or save the actual batch for debugging. In that case, you can temporarily store the tensors on the model during train_step:

1class DebugModel(tf.keras.Model):
2    def train_step(self, data):
3        x, y = data
4        self.last_batch_x = x
5        self.last_batch_y = y
6        return super().train_step(data)

Then a callback can read self.model.last_batch_x. This works, but it should be used carefully. Keeping full tensors around can increase memory use and may create confusion when running distributed training or graph-compiled code.

For long-running jobs, it is usually better to log compact summaries, indices, or a few sampled examples instead of storing whole batches.

Alternative approaches

There are other options, but they are usually heavier:

• write a custom training loop with tf.GradientTape
• wrap the dataset so each batch carries extra metadata
• emit debugging information with tf.print inside the model

A custom training loop gives total control and is the best fit when callbacks start fighting the design of your experiment.

Common Pitfalls

The most common mistake is expecting logs to contain x and y automatically. It will not.

Another problem is returning objects that are too large or not easily serializable in logs. Keep logs small and metric-like.

People also forget that tensors in callbacks may need conversion for display. Printing float(logs["loss"]) is clearer than dumping a raw TensorFlow tensor object.

Finally, if you override train_step, make sure you still update metrics and apply gradients correctly. A custom hook is not useful if it silently changes training behavior.

Summary

• Keras callbacks do not automatically expose raw batch tensors.
• The usual fix is to override train_step and return batch-derived values in logs.
• Full-batch access is possible by storing tensors on the model, but it increases complexity and memory usage.
• Use compact summaries unless you truly need the complete batch.
• If callback-based inspection feels awkward, switch to a custom training loop for full control.