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

fix checkpoint loading

Browse Source
This commit is contained in:
luni
2023-07-04 18:23:10 +02:00
parent 5ab1cb5fed
commit 8dc6e0c71a
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/checkpoints.py
+5 -3
View File
@@ -34,11 +34,13 @@ class CheckpointSaver:
def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=None, scaler: torch.nn.Module=None, iteration=''): 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 # 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')] checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')]
checkpoint = torch.load(checkpoints[-1]) iterations = np.asarray([int(f.split('_')[-1].split('.')[0]) for f in checkpoints])
iteration = checkpoints[-1].split('_')[-1].split('.')[0] last_iteration = np.argmax(iterations)
checkpoint = torch.load(checkpoints[last_iteration])
iteration = iterations[last_iteration]
model.load_state_dict(checkpoint['model']) model.load_state_dict(checkpoint['model'])
if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict()) 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']) 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}') print(f'Loaded model from {checkpoints[last_iteration]} at iteration {iteration}')
return int(iteration) if iteration != '' else 0 return int(iteration) if iteration != '' else 0