TensorBoard for visualising batch-level metrics.

tensorboard_batch_monitor.py

# TensorBoard for visualising batch-level metrics
# Based on code from various people at
# https://github.com/keras-team/keras/issues/6692

import tensorflow as tf
from keras.callbacks import TensorBoard

class TensorBoardBatchMonitor(TensorBoard):
    def __init__(self, log_every=1, **kwargs):
        super().__init__(**kwargs)
        self.log_every = log_every
        self.counter = 0
    
    def on_batch_end(self, batch, logs=None):
        self.counter += 1
        if self.counter % self.log_every == 0:
            for name, value in logs.items():
                if name in ['batch', 'size']:
                    continue
                summary = tf.Summary()
                summary_value = summary.value.add()
                summary_value.simple_value = value.item()
                summary_value.tag = name
                self.writer.add_summary(summary, self.counter)
            self.writer.flush()

    def on_epoch_end(self, epoch, logs=None):
        for name, value in logs.items():
            if name in ['acc', 'loss', 'batch', 'size']:
                continue
            summary = tf.Summary()
            summary_value = summary.value.add()
            summary_value.simple_value = value.item()
            summary_value.tag = name
            self.writer.add_summary(summary, self.counter)
        self.writer.flush()

If I include the on_epoch_end, I find the model no longer converges. If I just add pass to that code, then it converges but I don't get the validation metrics logged. Not sure what is happening.

@tonyreina: That's strange. on_epoch_end shouldn't have anything to do with whether the model converges or not. Is it not converging at all, or is it taking time to converge? If you increase log_every in the constructor, the logging will be done less frequently, and the training process will be a little faster. It might also be worth checking if on_epoch_end is being called at all, or how often it is being called, and what the items in logs are.