Tensorflow Keras Keep Loss of Every Batch

Question

I looking for advice/examples for the best practice in constructing my keras model architecture.

I've been fumbling with various iterations of the Model() subclass and the Functional model but have not been able to connect all the dots.

The model should have Custom Metrics and associated loss where-by: During training, metrics are calculated per batch and at the end of each epoch final metrics/loss are calculated from the mean of the batch calculations.

To my thinking I would need a Custom_Batch_Metric() that would maintain a list of the metrics for each batch, in addition to a Custom_Final_Metric() that would average the results from each batch. I'm not sure how to implement this.

For example... Per batch I would like to generate a metric and loss for the correlation of y_true, y_pred. At the end of my epoch I would like to then average (or find the max) of the batch correlations.

Would very much appreciate any reading someone could point me to regarding this sort of architecture.

Answer 1

A simple solution is to subclass tf.keras.callbacks.Callback and define the on_train_batch_end (or test). And then also on_epoch_end.

class SaveBatchLoss(tf.keras.callbacks.Callback):
    def on_train_batch_end(self, batch, logs=None): 
        batch_end_loss.append(logs['loss'])

import tensorflow as tf
from tensorflow.keras.layers import Dense
from sklearn.datasets import load_iris
import numpy as np

X, y = load_iris(return_X_y=True)
X = X.astype(np.float32)

ds = tf.data.Dataset.from_tensor_slices((X, y)).shuffle(25).batch(8)

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

model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', 
              metrics=['accuracy'])

batch_end_loss = list()

class SaveBatchLoss(tf.keras.callbacks.Callback):
    def on_train_batch_end(self, batch, logs=None):
        batch_end_loss.append(logs['loss'])

history = model.fit(ds, epochs=10, callbacks=SaveBatchLoss())

batch_end_loss[::20]

[1.2742226123809814,
 0.9069833755493164,
 0.9728888869285583,
 0.9536505937576294,
 0.8957988023757935,
 0.8624499440193176,
 0.7952101826667786,
 0.7765023112297058,
 0.7615134716033936,
 0.7278715968132019]

As a more complicated approach, append the loss values to a list at the end of every batch, and to another list at the end of each opoch. Like this:

train_loss_per_train_batch = list()
train_loss_per_train_epoch = 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())
    
    train_loss_per_train_epoch.append(train_loss.result())

A complete training script implementing this would be:

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)

class MyModel(tf.keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.d1 = tf.keras.layers.Dense(16, activation='relu')
        self.d2 = tf.keras.layers.Dense(32, activation='relu')
        self.d3 = tf.keras.layers.Dense(3, activation='softmax')

    def call(self, x, training=None, **kwargs):
        x = self.d1(x)
        x = self.d2(x)
        x = self.d3(x)
        return x

model = MyModel()

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()
train_loss_per_train_epoch = 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())

    train_loss_per_train_epoch.append(train_loss.result())

    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()))