python - How do I get non-aggregated batch statistics with tf.keras.callbacks?

Question

I am trying to train a model using the keras model.fit() method. This method returns a history object which contains loss values for each epoch - however I would like to have loss values for each individual batch.

Looking online I have found suggestions to use a custom callback class with an on_batch_end(self, logs={}) method. The problem is that this method only gets passed aggregated statistics that get reset each epoch. I would like to have individual statistics for each batch.

https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/Callback#on_train_batch_end

Answer 1

You could do that easily with a custom training loop, where you can just append a list with the loss value of every batch.

train_loss_per_train_batch.append(loss_value.numpy())

Here's how to do all of it:

import tensorflow as tf
import tensorflow_datasets as tfds

ds = tfds.load('iris', split='train', as_supervised=True)

train = ds.take(125).shuffle(16).batch(4)
test = ds.skip(125).take(25).shuffle(16).batch(4)

model = tf.keras.Sequential([
    tf.keras.layers.Dense(16, activation='relu'),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(3, activation='softmax')
])


loss_object = tf.losses.SparseCategoricalCrossentropy(from_logits=False)


def compute_loss(model, x, y, training):
  out = model(x, training=training)
  loss = loss_object(y_true=y, y_pred=out)
  return loss


def get_grad(model, x, y):
    with tf.GradientTape() as tape:
        loss = compute_loss(model, x, y, training=True)
    return loss, tape.gradient(loss, model.trainable_variables)


optimizer = tf.optimizers.Adam()

verbose = "Epoch {:2d} Loss: {:.3f} TLoss: {:.3f} Acc: {:.2%} TAcc: {:.2%}"

train_loss_per_train_batch = list()

for epoch in range(1, 25 + 1):
    train_loss = tf.metrics.Mean()
    train_acc = tf.metrics.SparseCategoricalAccuracy()
    test_loss = tf.metrics.Mean()
    test_acc = tf.metrics.SparseCategoricalAccuracy()

    for x, y in train:
        loss_value, grads = get_grad(model, x, y)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        train_loss.update_state(loss_value)
        train_acc.update_state(y, model(x, training=True))
        train_loss_per_train_batch.append(loss_value.numpy())

    for x, y in test:
        loss_value, _ = get_grad(model, x, y)
        test_loss.update_state(loss_value)
        test_acc.update_state(y, model(x, training=False))

    print(verbose.format(epoch,
                         train_loss.result(),
                         test_loss.result(),
                         train_acc.result(),
                         test_acc.result()))