From 8dc6e0c71a9e9a6b6655befad18ab0bde352278a Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 4 Jul 2023 18:23:10 +0200 Subject: [PATCH] fix checkpoint loading --- utils/checkpoints.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/utils/checkpoints.py b/utils/checkpoints.py index b564e40..8d23549 100644 --- a/utils/checkpoints.py +++ b/utils/checkpoints.py @@ -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=''): # 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] + 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']) 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}') + print(f'Loaded model from {checkpoints[last_iteration]} at iteration {iteration}') return int(iteration) if iteration != '' else 0