39 lines
1.7 KiB
Python
39 lines
1.7 KiB
Python
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_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': 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]
|