From 5db90e32ebd3f93b6af0386a90c4ca15e11e8849 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Wed, 23 Oct 2019 20:18:26 -0400 Subject: [PATCH] hpc restore takes priority over non hpc weights (#419) * hpc restore takes priority over non hpc weights * hpc restore takes priority over non hpc weights * hpc restore takes priority over non hpc weights * hpc restore takes priority over non hpc weights * hpc restore takes priority over non hpc weights * hpc restore takes priority over non hpc weights * hpc restore takes priority over non hpc weights --- pytorch_lightning/trainer/trainer_io.py | 29 +++++++++++++++---- ...ore_models.py => test_a_restore_models.py} | 2 +- tests/{test_logging.py => test_y_logging.py} | 0 3 files changed, 24 insertions(+), 7 deletions(-) rename tests/{test_restore_models.py => test_a_restore_models.py} (99%) rename tests/{test_logging.py => test_y_logging.py} (100%) diff --git a/pytorch_lightning/trainer/trainer_io.py b/pytorch_lightning/trainer/trainer_io.py index aa40e21d..5b483760 100644 --- a/pytorch_lightning/trainer/trainer_io.py +++ b/pytorch_lightning/trainer/trainer_io.py @@ -25,27 +25,37 @@ class TrainerIOMixin(object): def restore_weights(self, model): """ To restore weights we have two cases. - First, if we use the same experiment version, then restore the latest ckpt. - AFTER that, if we find weights from hpc checkpoint, then restore that. + First, attempt to restore hpc weights. If successful, don't restore + other weights. + + Otherwise, try to restore actual weights :param model: :return: """ - # restore weights if same exp version - self.restore_state_if_checkpoint_exists(model) # if script called from hpc resubmit, load weights - self.restore_hpc_weights_if_needed(model) + did_restore_hpc_weights = self.restore_hpc_weights_if_needed(model) + + if not did_restore_hpc_weights: + # restore weights if same exp version + self.restore_state_if_checkpoint_exists(model) # wait for all models to restore weights if self.use_ddp or self.use_ddp2: # wait for all processes to catch up dist.barrier() + # clear cache after restore + if self.on_gpu: + torch.cuda.empty_cache() + def restore_state_if_checkpoint_exists(self, model): + did_restore = False + # do nothing if there's not dir or callback no_ckpt_callback = (self.checkpoint_callback is None) or (not self.checkpoint_callback) if no_ckpt_callback or not os.path.exists(self.checkpoint_callback.filepath): - return + return did_restore # restore trainer state and model if there is a weight for this experiment last_epoch = -1 @@ -71,6 +81,9 @@ class TrainerIOMixin(object): last_ckpt_path = os.path.join(self.checkpoint_callback.filepath, last_ckpt_name) self.restore(last_ckpt_path, self.on_gpu) print(f'model and trainer restored from checkpoint: {last_ckpt_path}') + did_restore = True + + return did_restore # -------------------- # HPC SIGNAL HANDLING @@ -198,6 +211,8 @@ class TrainerIOMixin(object): :param model: :return: """ + did_restore = False + # look for hpc weights folderpath = self.weights_save_path if os.path.exists(folderpath): @@ -207,6 +222,8 @@ class TrainerIOMixin(object): # if hpc weights exist restore model if len(hpc_weight_paths) > 0: self.hpc_load(folderpath, self.on_gpu) + did_restore = True + return did_restore def restore_training_state(self, checkpoint): """ diff --git a/tests/test_restore_models.py b/tests/test_a_restore_models.py similarity index 99% rename from tests/test_restore_models.py rename to tests/test_a_restore_models.py index afa21cfd..5b747aec 100644 --- a/tests/test_restore_models.py +++ b/tests/test_a_restore_models.py @@ -327,8 +327,8 @@ def test_cpu_restore_training(): # set the epoch start hook so we can predict before the model does the full training def assert_good_acc(): - assert trainer.current_epoch > 0 assert trainer.current_epoch == real_global_epoch + assert trainer.current_epoch >= 0 # if model and state loaded correctly, predictions will be good even though we # haven't trained with the new loaded model diff --git a/tests/test_logging.py b/tests/test_y_logging.py similarity index 100% rename from tests/test_logging.py rename to tests/test_y_logging.py