new scripts

utils/checkpoints.py

@@ -33,12 +33,16 @@ class CheckpointSaver:
 
 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')]
+    if iteration == '':
+        checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith('.pt')]
+    else:
+        checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'iteration_{iteration}.pt')]
+        assert len(checkpoints) == 1
     iterations = np.asarray([int(f.split('_')[-1].split('.')[0]) for f in checkpoints])
     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'], strict=False)
     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[last_iteration]} at iteration {iteration}')