45 lines
2.0 KiB
Python
45 lines
2.0 KiB
Python
import torch
|
|
import numpy as np
|
|
import os
|
|
import logging
|
|
|
|
|
|
class CheckpointSaver:
|
|
def __init__(self, dirpath, 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.top_model_paths = []
|
|
|
|
def __call__(self, model, iteration, optimizer, scaler):
|
|
model_path = os.path.join(self.dirpath, f'iteration_{iteration}.pt')
|
|
torch.save({ 'model': model.state_dict(), 'optimizer': optimizer, 'scaler': scaler }, model_path)
|
|
self.top_model_paths.append({'path': model_path})
|
|
if len(self.top_model_paths)>self.top_n:
|
|
self.cleanup()
|
|
|
|
def cleanup(self):
|
|
to_remove = self.top_model_paths[:len(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[len(self.top_model_paths)-self.top_n:]
|
|
|
|
|
|
def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=None, scaler: torch.nn.Module=None, iteration=''):
|
|
# list of checkpoint files in the folder checkpoints
|
|
checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')]
|
|
checkpoint = torch.load(checkpoints[-1])
|
|
iteration = checkpoints[-1].split('_')[-1].split('.')[0]
|
|
model.load_state_dict(checkpoint['model'])
|
|
if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict())
|
|
if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler'])
|
|
print(f'Loaded model from {checkpoints[-1]} at iteration {iteration}')
|
|
return int(iteration) if iteration != '' else 0
|
|
|