test_on_batch/train_on_batch returns cumulative metric, causing inconsistent loss after load_model

[coderW52]

When using train_on_batch for training and test_on_batch for validation, the reported loss is not deterministic after loading a saved model.
Even though model weights and optimizer state are correctly restored, test_on_batch returns a cumulative metric state, which can lead to misleading loss values in custom training loops.

import numpy as np
import keras
import tensorflow as tf

### --- Minimal model ---
model = keras.Sequential([
    keras.layers.LSTM(32, input_shape=(10, 4)),
    keras.layers.Dense(1)
])
model.compile(optimizer='adam', loss='mse')

### --- Dummy data ---
X = np.random.randn(100, 10, 4).astype("float32")
y = np.random.randn(100, 1).astype("float32")

X_batch = X[:32]
y_batch = y[:32]

### --- Train for 5 batches to accumulate metric state ---
for _ in range(5):
    model.train_on_batch(X_batch, y_batch)

### --- Save the model ---
model.save("bug_repro.keras")

### --- BEFORE load_model ---
loss_before = model.test_on_batch(X_batch, y_batch)
y_pred_before = model(X_batch, training=False)
mse_before = float(tf.reduce_mean(tf.square(y_pred_before - y_batch)))

print(f"BEFORE load_model | test_on_batch : {loss_before:.6f}")
print(f"BEFORE load_model | manual MSE    : {mse_before:.6f}")

### --- load_model creates a new model object with a fresh metric state ---
loaded_model = keras.models.load_model("bug_repro.keras")

### --- AFTER load_model ---
loss_after = loaded_model.test_on_batch(X_batch, y_batch)
y_pred_after = loaded_model(X_batch, training=False)
mse_after = float(tf.reduce_mean(tf.square(y_pred_after - y_batch)))

print(f"\nAFTER  load_model | test_on_batch : {loss_after:.6f}")
print(f"AFTER  load_model | manual MSE    : {mse_after:.6f}")

### --- Diagnosis ---
print(f"\n{'='*50}")
print(f"Weights identical          : {all(np.allclose(w1.numpy(), w2.numpy()) for w1, w2 in zip(model.weights, loaded_model.weights))}")
print(f"Manual MSE identical       : {abs(mse_before - mse_after) < 1e-6}")
print(f"test_on_batch identical    : {abs(loss_before - loss_after) < 1e-6}")

#############################################################
BEFORE load_model | test_on_batch : 1.608474
BEFORE load_model | manual MSE    : 1.573039

AFTER  load_model | test_on_batch : 1.573039
AFTER  load_model | manual MSE    : 1.573039

==================================================
Weights identical          : True
Manual MSE identical       : True
test_on_batch identical    : False

Expected behavior: test_on_batch should return the same value as manual MSE

Actual behavior: test_on_batch returns a cumulative average from the accumulated metric state, which is reset upon load_model().
This causes inconsistent loss reporting when resuming training after load_model(), even though model weights
and optimizer state are correctly restored.

Environment

• Keras version: 3.13.0
• TensorFlow version: 2.19.0
• Python version: 3.11
• GPU: Yes, cuDNN 9.3

[github-actions]

Thank you for providing a detailed reproduction script.

In Keras 3, train_on_batch and test_on_batch are designed to return the current state of the model's metrics. This includes the loss, which is tracked by an internal Mean metric (the _loss_tracker). These metrics accumulate their state across multiple calls until model.reset_metrics() is invoked.

When you use keras.models.load_model(), a new model instance is created. While weights and optimizer states are restored, the volatile metric states (like the batch count and sum in the loss tracker) are initialized to zero. This explains why test_on_batch returns the non-cumulative batch loss immediately after loading, but returns a cumulative average after several training steps.

Workaround

If you want test_on_batch to return the loss for only the current batch (matching your manual MSE calculation), you can call model.reset_metrics() before calling test_on_batch:

model.reset_metrics()
loss = model.test_on_batch(X_batch, y_batch)

We will evaluate if this behavior should be more explicitly documented or if metric states should be optionally preservable during serialization. Tagging for team review.

[coderW52]

Thank you for the clarification. I understand that the cumulative metric behavior in train_on_batch and test_on_batch is by design, and that metric states are not serialized with load_model(). For my use case—resuming training from a checkpoint while monitoring validation loss—this makes batch-level loss reporting unstable. I will calculate epoch-level losses manually for logging and checkpointing. I appreciate your time and hope future updates may provide optional metric-state serialization for custom loops.

[amitsrivastava78]

 train_on_batch/test_on_batch return a cumulative running average (sum/count) tracked by an internal Mean metric (_loss_tracker). After 5 training batches, the 6th test_on_batch call returns the average of all 6 — not just the current batch. load_model resets this incidentally by creating a fresh model instance with zeroed metric state.
Fix: Call model.reset_metrics() before test_on_batch to get the true per-batch loss:

model.reset_metrics()
loss = model.test_on_batch(X_batch, y_batch)

Not a bug — it's a documentation gap