luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

upscaling pipeline

Browse Source
This commit is contained in:
luni
2023-05-21 13:52:42 +00:00
parent fc27a3cedf
commit 75df8a09cc
5 changed files with 322 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/checkpoints.py
+4 -3
View File
@@ -16,9 +16,9 @@ class CheckpointSaver:
self.top_n = top_n
self.top_model_paths = []
def __call__(self, model, iteration, optimizer):
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 }, model_path)
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()
@@ -31,12 +31,13 @@ class CheckpointSaver:
self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:]
def load_model(model, checkpoint, optimizer=None, iteration=''):
def load_model(model, checkpoint, optimizer=None, scaler=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'])
if scaler is not None: scaler.load_state_dict(checkpoint['scaler'])
return int(iteration) if iteration != '' else 0