66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
import os
|
|
|
|
import torch
|
|
|
|
|
|
class Checkpoint:
|
|
"""
|
|
Handles saving and loading checkpoints.
|
|
|
|
Args:
|
|
checkpoint_dir (str): path where checkpoints are saved
|
|
device: PyTorch device onto which loaded weights should be mapped
|
|
kwargs: PyTorch modules whose state should be checkpointed
|
|
"""
|
|
|
|
def __init__(self, checkpoint_dir="./chkpts", device=None, **kwargs):
|
|
self.module_dict = kwargs
|
|
self.device = device
|
|
# self.checkpoint_dir = checkpoint_dir
|
|
# if not os.path.exists(checkpoint_dir):
|
|
# os.makedirs(checkpoint_dir)
|
|
|
|
# def save(self, filename, **kwargs):
|
|
# """Saves the current module states
|
|
# Args:
|
|
# filename (str): name of checkpoint file
|
|
# kwargs: Additional state to save
|
|
# """
|
|
# if not os.path.isabs(filename):
|
|
# filename = os.path.join(self.checkpoint_dir, filename)
|
|
|
|
# outdict = kwargs
|
|
# for k, v in self.module_dict.items():
|
|
# if k in outdict:
|
|
# print(
|
|
# f"Warning: Checkpoint key '{k}' overloaded. Defaulting to saving state_dict {v}."
|
|
# )
|
|
# if v is not None:
|
|
# outdict[k] = v.state_dict()
|
|
# torch.save(outdict, filename)
|
|
|
|
def load(self, filename):
|
|
"""Loads a checkpoint from file.
|
|
Args:
|
|
filename (str): Name of checkpoint file.
|
|
Returns:
|
|
Dictionary containing checkpoint data which does not correspond to registered modules.
|
|
"""
|
|
|
|
# if not os.path.isabs(filename):
|
|
# filename = os.path.join(self.checkpoint_dir, filename)
|
|
|
|
print(f"Loading checkpoint from file {filename}...")
|
|
state_dict = torch.load(filename, map_location=self.device)
|
|
|
|
for k, v in self.module_dict.items():
|
|
if k in state_dict:
|
|
v.load_state_dict(state_dict[k])
|
|
else:
|
|
print(f'Warning: Could not find "{k}" in checkpoint!')
|
|
|
|
remaining_state = {
|
|
k: v for k, v in state_dict.items() if k not in self.module_dict
|
|
}
|
|
return remaining_state
|