PlotLossesCallback.py

import datetime
import tensorflow as tf
import numpy as np
from matplotlib import pyplot as plt
from IPython.display import clear_output

class TrainTimeCallback(tf.keras.callbacks.Callback):
    def on_epoch_begin(self, epoch, logs=None):
        self.epoch_start_time = datetime.datetime.now()

    def on_epoch_end(self, epoch, logs=None):
        logging.info('\nepoch train time: (hh:mm:ss.ms) {}'.format(
            datetime.datetime.now() - self.epoch_start_time))
        if not self.params is None:
            if 'steps' in self.params and self.params['steps']:
                epoch_milliseconds = (datetime.datetime.now(
                ) - self.epoch_start_time).total_seconds() * 1000
                logging.info(
                    '{} ms/step'.format(epoch_milliseconds / self.params['steps']))
                if BATCH_SIZE is not None:
                    logging.info('{} microseconds/example'.format(
                        1000 * epoch_milliseconds / self.params['steps'] / BATCH_SIZE))
                    


class PlotLossesCallback(tf.keras.callbacks.Callback):
    def on_train_begin(self, logs={}):
        self.i = 0
        self.x = []
        self.losses = []
        self.val_losses = []
        self.fig = plt.figure()
        self.logs = []
        self.learning_rates = []

    def on_epoch_end(self, epoch, logs={}):
        
        self.logs.append(logs)
        self.x.append(self.i)
        self.losses.append(logs.get('loss'))
        self.val_losses.append(logs.get('val_loss'))
        self.learning_rates.append(tf.keras.backend.get_value(self.model.optimizer.lr))
        self.i += 1
        
        clear_output(wait=True)
        
        fig, axs = plt.subplots(1,2,figsize=(12,6))
        fig.suptitle('Vertically stacked subplots')
        losses_plt = axs[0]
        losses_plt.plot(self.x, self.losses, label="loss")
        losses_plt.plot(self.x, self.val_losses, label="val_loss")
        losses_plt.legend()
        
        lr_plt = axs[1]
        lr_plt.plot(self.x, self.learning_rates, label="learning rate")
        lr_plt.legend()
        
        plt.show()