Implementation of early stopping and checkpointing in PyTorch

EarlyStoppingCheckpoint.py

import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import save_image
from tqdm import tqdm
from collections import defaultdict

class EarlyStoppingCheckpoint:
    def __init__(self, model, save_path, metric_name, mode='max', patience=10):
        self.model = model
        self.save_path = save_path
        self.metric_name = metric_name
        self.mode = mode
        self.patience = patience
        self.best_metric = None
        self.best_epoch = None
        self.epochs_since_improvement = 0

        if self.mode == 'max':
            self.best_metric = float('-inf')
        elif self.mode == 'min':
            self.best_metric = float('inf')
        else:
            raise ValueError(f"Invalid mode: {self.mode}. Must be 'max' or 'min'.")

    def __call__(self, metric, epoch):
        if self.mode == 'max' and metric > self.best_metric:
            self.best_metric = metric
            self.best_epoch = epoch
            self.epochs_since_improvement = 0
            self.save_checkpoint()
        elif self.mode == 'min' and metric < self.best_metric:
            self.best_metric = metric
            self.best_epoch = epoch
            self.epochs_since_improvement = 0
            self.save_checkpoint()
        else:
            self.epochs_since_improvement += 1
            if self.epochs_since_improvement >= self.patience:
                print(f"Stopping early. No improvement in {self.patience} epochs.")
                return True  # Early stopping
        return False  # Continue training

    def save_checkpoint(self):
        checkpoint = {'epoch': self.best_epoch, 'model': self.model.state_dict(), 'metric': self.best_metric}
        torch.save(checkpoint, self.save_path)
        print(f"Saved checkpoint at epoch {self.best_epoch}, with {self.metric_name} of {self.best_metric:.4f}.")