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

@@ -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):
         """

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