callbacks.py

import math
import json
import pickle


class Callback:
    def __init__():
        pass

    def on_epoch_end(log_train, log_valid, model):
        pass


class MetricsLogger(Callback):
    def __init__(self, log_dest):
        self.history = {
            'train': [],
            'valid': []
        }
        self.log_dest = log_dest

    def on_epoch_end(self, log_train, log_valid, model):
        log_train['epoch'] = model.epoch
        log_valid['epoch'] = model.epoch
        self.history['train'].append(log_train)
        self.history['valid'].append(log_valid)
        with open(self.log_dest, 'w') as f:
            json.dump(self.history, f, indent='    ')
            
    def load(self, epoch):
        with open(self.log_dest, 'r') as f:
            self.history = json.loads(f.read())
        self.history['train'] = self.history['train'][:epoch + 1]
        self.history['valid'] = self.history['valid'][:epoch + 1]

class ModelCheckpoint(Callback):
    def __init__(self, filepath,
                 monitor='loss',
                 verbose=0,
                 mode='min',
                 early_stop=10):
        self._filepath = filepath
        self._verbose = verbose
        self._monitor = monitor
        self._best = math.inf if mode == 'min' else - math.inf
        self._mode = mode
        self.early_stop = early_stop
        self.early_stop_counter = -1

    def on_epoch_end(self, log_train, log_valid, model):
        score = log_valid[self._monitor]
        if self._mode == 'min':
            if score < self._best:
                self.early_stop_counter = -1
                self._best = score
                model.save(self._filepath)
                if self._verbose > 0:
                    print('Best model saved (%f)' % score)

        elif self._mode == 'max':
            if score > self._best:
                self.early_stop_counter = -1
                self._best = score
                model.save('{}'.format(self._filepath, model.epoch)) # save only best model ignore epoch
                if self._verbose > 0:
                    print('Best model saved (%f)' % score)

        elif self._mode == 'all':
            self.early_stop_counter = -1
            model.save('{}.{}'.format(self._filepath, model.epoch))
        
        self.early_stop_counter += 1
        if self.early_stop_counter > self.early_stop:
            exit()