How to save all your trained model weights locally after every epoch

Introduction

Checkpoint Saver in PyTorch

We need to be able to provide the user the option to select a certain number of models that users want to save.
We need to create a valid name for the model for ease of access later on.
Give users the ability to provide a directory where they want to save the model weights.
Either save model weights based on highest validation metric scores or lowest validation loss.

import numpy as np 
import os
import logging

class CheckpointSaver:
    def __init__(self, dirpath, decreasing=True, top_n=5):
        """
        dirpath: Directory path where to store all model weights 
        decreasing: If decreasing is `True`, then lower metric is better
        top_n: Total number of models to track based on validation metric value
        """
        if not os.path.exists(dirpath): os.makedirs(dirpath)
        self.dirpath = dirpath
        self.top_n = top_n 
        self.decreasing = decreasing
        self.top_model_paths = []
        self.best_metric_val = np.Inf if decreasing else -np.Inf
        
    def __call__(self, model, epoch, metric_val):
        model_path = os.path.join(self.dirpath, model.__class__.__name__ + f'_epoch{epoch}.pt')
        save = metric_val<self.best_metric_val if self.decreasing else metric_val>self.best_metric_val
        if save: 
            logging.info(f"Current metric value better than {metric_val} better than best {self.best_metric_val}, saving model at {model_path}")
            self.best_metric_val = metric_val
            torch.save(model.state_dict(), model_path)
            self.top_model_paths.append({'path': model_path, 'score': metric_val})
            self.top_model_paths = sorted(self.top_model_paths, key=lambda o: o['score'], reverse=not self.decreasing)
        if len(self.top_model_paths)>self.top_n: 
            self.cleanup()
    
    def cleanup(self):
        to_remove = self.top_model_paths[self.top_n:]
        logging.info(f"Removing extra models.. {to_remove}")
        for o in to_remove:
            os.remove(o['path'])
        self.top_model_paths = self.top_model_paths[:self.top_n]

How does it work?

top_n : When we call the checkpoint saver, inside __call__ method, the condition if len(self.top_model_paths)>self.top_n checks if the total number of models saved is more than top_n, in that case, the checkpoint saver cleans the extra models that have a lower performance.
The model_path inside the checkpoint saver takes in the model.__class__.__name__ and also the epoch number, thus providing a valid name to save the model.
All models get stored inside dirpath
If we are saving models where decreasing metric value is better, that is we are using validation loss to track models, in that case we set decreasing=True. When decreasing is True, then the best validation value in the beginning is np.Inf. Therefore, at every epoch we check if the new metric value is less than the previous value. If it is, then we save the model. Otherwise, we do not save the model. Reverse logic follows for when decreasing=False. We only save the model if the current score is better than the previous epoch's score.

A working example

Summary of saving models using Checkpoint Saver

Storage: Each model weights are about 100 MBs each, so we will need a few GBs to store the model weights. Generally, it takes a lot of space if you are running a lot of experiments! Also, bigger models take more space.
Model names and scores: Sure, we get a list of models with model names and epochs as part of the names, but how do we know which experiments were these models a part of? What if we accidentally overwrite the model weights? How do we know what was the score for Resnet_epoch0 model? One quick way could be to add metric scores to model names as well but that would make it really messy. For example, each model would be called something like Resnet_epoch0_0.2313, and IMHO this is definitely not the best way to go about organizing your experiments.
Store model names in Excel sheet with Experiment names: We would have to manually keep track of experiments that we run and model weights stored for those experiments.
Which model is better? Since model names have only epoch information associated with them, how do we know if the model weights for epoch 4 were better or that for epoch 3?
How does model performance track overtime? This way of saving model weights and progress shows us no information about how model performance has changed over time per epoch.

Experiment Tracking with Weights and Biases

Figure-1: Sample dashboard on Melanoma project when using W&B for experiment tracking

A more scalable way of saving your model weights

What are W&B Artifacts?

W&B Artifacts was designed to make it effortless to version your datasets and models, regardless of whether you want to store your files with W&B or whether you already have a bucket you want W&B to track. Once you've tracked your datasets or model files, W&B will automatically log each and every modification, giving you a complete and auditable history of changes to your files.

wandb.init()
artifact = wandb.Artifact(<enter_filename>, type='model')
artifact.add_file(<file_path>)
wandb.run.log_artifact(artifact)

W&B automatically versions each of the model artifacts for you.
W&B does not overwrite model artifacts: Even if you store two artifacts with the same name, unless you specify to update a particular version of the model artifact, W&B will create a new model version for you. Please note that this is not possible when you are storing model weights locally as we overwrite files if they have the same name.
All part of a bigger ecosystem: By using artifacts and adding them to experiment tracking, you become part of a bigger W&B ecosystem and it means that all your relevant data is in one place. So no longer do you need to maintain an excel sheet with the experiments that you run and file paths.

class CheckpointSaver:
    def __init__(self, dirpath, decreasing=True, top_n=5):
        """
        dirpath: Directory path where to store all model weights 
        decreasing: If decreasing is `True`, then lower metric is better
        top_n: Total number of models to track based on validation metric value
        """
        if not os.path.exists(dirpath): os.makedirs(dirpath)
        self.dirpath = dirpath
        self.top_n = top_n 
        self.decreasing = decreasing
        self.top_model_paths = []
        self.best_metric_val = np.Inf if decreasing else -np.Inf
        
    def __call__(self, model, epoch, metric_val):
        model_path = os.path.join(self.dirpath, model.__class__.__name__ + f'_epoch{epoch}.pt')
        save = metric_val<self.best_metric_val if self.decreasing else metric_val>self.best_metric_val
        if save: 
            logging.info(f"Current metric value better than {metric_val} better than best {self.best_metric_val}, saving model at {model_path}, & logging model weights to W&B.")
            self.best_metric_val = metric_val
            torch.save(model.state_dict(), model_path)
            self.log_artifact(f'model-ckpt-epoch-{epoch}.pt', model_path, metric_val)
            self.top_model_paths.append({'path': model_path, 'score': metric_val})
            self.top_model_paths = sorted(self.top_model_paths, key=lambda o: o['score'], reverse=not self.decreasing)
        if len(self.top_model_paths)>self.top_n: 
            self.cleanup()
    
    def log_artifact(self, filename, model_path, metric_val):
        artifact = wandb.Artifact(filename, type='model', metadata={'Validation score': metric_val})
        artifact.add_file(model_path)
        wandb.run.log_artifact(artifact)        
    
    def cleanup(self):
        to_remove = self.top_model_paths[self.top_n:]
        logging.info(f"Removing extra models.. {to_remove}")
        for o in to_remove:
            os.remove(o['path'])
        self.top_model_paths = self.top_model_paths[:self.top_n]

Do you know the best part about using W&B artifacts? I can also add metadata! When we were saving model files locally, we had no idea about the validation scores

Conclusion