diff --git a/utils/checkpoints.py b/utils/checkpoints.py new file mode 100644 index 0000000..8ba4c5c --- /dev/null +++ b/utils/checkpoints.py @@ -0,0 +1,38 @@ +import torch +import numpy as np +import os +import logging + + +class CheckpointSaver: + def __init__(self, dirpath, decreasing=False, 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, optimizer, metric_val=np.Inf): + model_path = os.path.join(self.dirpath, model.__class__.__name__ + f'_epoch{epoch}.pt') + save = metric_valself.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': model.state_dict(), 'optimizer':optimizer }, 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]