earlystoping.py

import torch

class Earlystopping:
    def __init__(self, number=50, path='../ssd/checkpoint.pt'):
        self.number = number
        self.path = path
        self.counter = 0
        self.epoch_count = 0
        self.best_epoch_num = 1
        self.max_acc = None
        self.stop_now = False

    def __call__(self, model, acc):
        if self.max_acc is None:
            self.epoch_count += 1
            self.max_acc = acc
            self.save_model(model)
        elif acc > self.max_acc:
            self.epoch_count += 1
            self.best_epoch_num = self.epoch_count
            print('New maximum accuracy: {:.4f}% -> {:.4f}%\n'.format(self.max_acc, acc))
            self.max_acc = acc
            self.save_model(model)
            self.counter = 0
        else:
            self.epoch_count += 1
            print('Current maximum accuracy: {:.4f}%\n'
                  'Early stopping counter: {:d}/{:d}\n'.format(self.max_acc, self.counter, self.number))
            if self.counter >= self.number:
                self.stop_now = True
            self.counter += 1

    def save_model(self, model):
        torch.save(model.state_dict(), self.path)