From d3f19c83210941da0e4796f0311137881d82a5ed Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 06:55:05 -0400
Subject: [PATCH 01/22] added auto restore

---
 pytorch_lightning/models/trainer.py           | 24 ++++++++++++++++++
 pytorch_lightning/root_module/model_saving.py | 25 ++++++++++++++++---
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index 2655809c..c9c4d850 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -244,6 +244,30 @@ class Trainer(TrainerIO):
             '''
             raise ModuleNotFoundError(msg)
 
+        # restore training and model
+        self.restore_state_if_existing_checkpoint()
+
+    def restore_state_if_existing_checkpoint(self):
+        # restore trainer state and model if there is a weight for this experiment
+        last_epoch = -1
+        last_ckpt_name = None
+
+        # find last epoch
+        checkpoints = os.listdir(self.checkpoint_callback.filepath)
+        for name in checkpoints:
+            if '.ckpt' in name:
+                epoch = name.split('epoch_')[1]
+                epoch = re.sub('[^0-9]', '' ,epoch)
+
+                if epoch > last_epoch:
+                    last_epoch = epoch
+                    last_ckpt_name = name
+
+        # restore last checkpoint
+        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}')
+
     @property
     def data_parallel(self):
         return self.use_dp or self.use_ddp
diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py
index 0bde0943..d9c2e3a5 100644
--- a/pytorch_lightning/root_module/model_saving.py
+++ b/pytorch_lightning/root_module/model_saving.py
@@ -58,6 +58,25 @@ class TrainerIO(object):
         # do the actual save
         torch.save(checkpoint, filepath)
 
+    def restore(self, checkpoint_path, on_gpu):
+
+        if on_gpu:
+            checkpoint = torch.load(checkpoint_path)
+        else:
+            checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
+
+        # load training state (affects trainer only)
+        self.restore_training_state(checkpoint)
+
+        # load model state
+        model = self.__get_model()
+
+        # load the state_dict on the model automatically
+        model.load_state_dict(checkpoint['state_dict'])
+
+        # call model hook
+        model.on_hpc_load(checkpoint)
+
     def dump_checkpoint(self):
 
         checkpoint = {
@@ -198,15 +217,15 @@ class TrainerIO(object):
         # call model hook
         model.on_hpc_load(checkpoint)
 
-    def max_ckpt_in_folder(self, path):
+    def max_ckpt_in_folder(self, path, name_key='ckpt_'):
         files = os.listdir(path)
-        files = [x for x in files if 'ckpt_' in x]
+        files = [x for x in files if name_key in x]
         if len(files) == 0:
             return 0
 
         ckpt_vs = []
         for name in files:
-            name = name.split('ckpt_')[-1]
+            name = name.split(name_key)[-1]
             name = re.sub('[^0-9]', '', name)
             ckpt_vs.append(int(name))
 

From 47a691f1583c317b497d389d68a4bd7b2cc814d3 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:09:37 -0400
Subject: [PATCH 02/22] updated tests and docs

---
 README.md                     |  1 +
 docs/Trainer/Checkpointing.md | 17 ++++++++
 docs/Trainer/index.md         |  1 +
 docs/index.md                 |  1 +
 tests/test_models.py          | 75 ++++++++++++++++++++++++++++++++++-
 5 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 7542eeef..76ced232 100644
--- a/README.md
+++ b/README.md
@@ -259,6 +259,7 @@ tensorboard --logdir /some/path
 
 - [Model saving](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)
 - [Model loading](https://williamfalcon.github.io/pytorch-lightning/LightningModule/methods/#load-from-metrics) 
+- [Restoring training session](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#restoring-training-session)
 
 ###### Computing cluster (SLURM)    
 
diff --git a/docs/Trainer/Checkpointing.md b/docs/Trainer/Checkpointing.md
index db0d7d50..b0a5281d 100644
--- a/docs/Trainer/Checkpointing.md
+++ b/docs/Trainer/Checkpointing.md
@@ -18,5 +18,22 @@ checkpoint_callback = ModelCheckpoint(
 trainer = Trainer(checkpoint_callback=checkpoint_callback)
 ```
 
+---
+### Restoring training session 
+You might want to not only load a model but also continue training it. Use this method to
+restore the trainer state as well. This will continue from the epoch and global step you last left off.  
+However, the dataloaders will start from the first batch again (if you shuffled it shouldn't matter).   
+
+Lightning will restore the session if you pass an experiment with the same version and there's a saved checkpoint.   
+``` {.python}
+from test_tube import Experiment
+
+exp = Experiment(version=a_previous_version_with_a_saved_checkpoint)
+Trainer(experiment=exp)
+
+trainer = Trainer(checkpoint_callback=checkpoint_callback)
+# the trainer is now restored
+```
+
 
 
diff --git a/docs/Trainer/index.md b/docs/Trainer/index.md
index 88d85abe..7d363a21 100644
--- a/docs/Trainer/index.md
+++ b/docs/Trainer/index.md
@@ -21,6 +21,7 @@ But of course the fun is in all the advanced things it can do:
 
 - [Model saving](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)
 - [Model loading](https://williamfalcon.github.io/pytorch-lightning/LightningModule/methods/#load-from-metrics) 
+- [Restoring training session](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#restoring-training-session)
 
 **Computing cluster (SLURM)**    
 
diff --git a/docs/index.md b/docs/index.md
index 0107897d..45f3e25c 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -28,6 +28,7 @@ one could be a seq-2-seq model, both (optionally) ran by the same trainer file.
 
 - [Model saving](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)
 - [Model loading](https://williamfalcon.github.io/pytorch-lightning/LightningModule/methods/#load-from-metrics) 
+- [Restoring training session](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#restoring-training-session)
 
 ###### Computing cluster (SLURM)    
 
diff --git a/tests/test_models.py b/tests/test_models.py
index 73ba2e43..1bc6b8dd 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -52,6 +52,77 @@ def test_amp_gpu_ddp():
     run_gpu_model_test(trainer_options, model, hparams)
 
 
+def test_cpu_restore_training():
+    """
+    Verify continue training session on CPU
+    :return:
+    """
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    test_exp_version = 10
+    exp = get_exp(False, version=test_exp_version)
+    exp.argparse(hparams)
+    exp.save()
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    real_global_step = trainer.global_step
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # predict with trained model before saving
+    # make a prediction
+    for batch in model.test_dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    model.eval()
+    pred_before_saving = model(x)
+
+    # wipe-out trainer and model
+    # retrain with not much data... this simulates picking training back up after slurm
+    # we want to see if the weights come back correctly
+    new_exp = get_exp(False, version=test_exp_version)
+    trainer_options = dict(
+        max_nb_epochs=1,
+        experiment=new_exp,
+        checkpoint_callback=ModelCheckpoint(save_dir),
+    )
+    trainer = Trainer(**trainer_options)
+    model = LightningTestModel(hparams)
+
+    # set the epoch start hook so we can predict before the model does the full training
+    def assert_pred_same():
+        assert trainer.global_step == real_global_step and trainer.global_step > 0
+
+        # predict with loaded model to make sure answers are the same
+        trainer.model.eval()
+        new_pred = trainer.model(x)
+        assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
+
+    model.on_epoch_start = assert_pred_same
+
+    # by calling fit again, we trigger training, loading weights from the cluster
+    # and our hook to predict using current model before any more weight updates
+    trainer.fit(model)
+
+    clear_save_dir()
+
+
 def test_cpu_slurm_save_load():
     """
     Verify model save/load/checkpoint on CPU
@@ -610,10 +681,10 @@ def get_model():
     return model, hparams
 
 
-def get_exp(debug=True):
+def get_exp(debug=True, version=None):
     # set up exp object without actually saving logs
     root_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
+    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir', version=version)
     return exp
 
 

From a895bf1b712861ffc806001ab56c9bbfbacbbc8f Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:15:23 -0400
Subject: [PATCH 03/22] fixed none name

---
 pytorch_lightning/models/trainer.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index c9c4d850..fc1fb8f5 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -264,9 +264,10 @@ class Trainer(TrainerIO):
                     last_ckpt_name = name
 
         # restore last checkpoint
-        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}')
+        if last_ckpt_name is not None:
+            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}')
 
     @property
     def data_parallel(self):

From 82d63a9677d58595b080a7ff9a5acd48aff2287e Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:25:20 -0400
Subject: [PATCH 04/22] fixed none name

---
 pytorch_lightning/models/trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index fc1fb8f5..4d72618b 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -257,7 +257,7 @@ class Trainer(TrainerIO):
         for name in checkpoints:
             if '.ckpt' in name:
                 epoch = name.split('epoch_')[1]
-                epoch = re.sub('[^0-9]', '' ,epoch)
+                epoch = int(re.sub('[^0-9]', '' ,epoch))
 
                 if epoch > last_epoch:
                     last_epoch = epoch

From 95ec072d1e5b323aeaf230cb85955429d7319018 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:30:02 -0400
Subject: [PATCH 05/22] removed bad hook call

---
 pytorch_lightning/root_module/model_saving.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py
index d9c2e3a5..00b95460 100644
--- a/pytorch_lightning/root_module/model_saving.py
+++ b/pytorch_lightning/root_module/model_saving.py
@@ -74,9 +74,6 @@ class TrainerIO(object):
         # load the state_dict on the model automatically
         model.load_state_dict(checkpoint['state_dict'])
 
-        # call model hook
-        model.on_hpc_load(checkpoint)
-
     def dump_checkpoint(self):
 
         checkpoint = {

From 2575b157a474049100e74426c7ac4d92069d8b73 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:32:33 -0400
Subject: [PATCH 06/22] removed bad hook call

---
 tests/test_models.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/test_models.py b/tests/test_models.py
index 1bc6b8dd..509d529c 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -70,6 +70,9 @@ def test_cpu_restore_training():
 
     trainer_options = dict(
         max_nb_epochs=1,
+        val_check_interval=0.50,
+        val_percent_check=0.2,
+        train_percent_check=0.2,
         experiment=exp,
         checkpoint_callback=ModelCheckpoint(save_dir)
     )

From 9713c41bf419e1b3858880198cbf90a379e0d1c0 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:33:08 -0400
Subject: [PATCH 07/22] removed bad hook call

---
 tests/test_models.py | 57 +++++++++++++++++++++++---------------------
 1 file changed, 30 insertions(+), 27 deletions(-)

diff --git a/tests/test_models.py b/tests/test_models.py
index 509d529c..1708134a 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -24,33 +24,6 @@ np.random.seed(SEED)
 # ------------------------------------------------------------------------
 # TESTS
 # ------------------------------------------------------------------------
-def test_amp_gpu_ddp():
-    """
-    Make sure DDP + AMP work
-    :return:
-    """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
-        return
-
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-
-    hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    trainer_options = dict(
-        progress_bar=True,
-        max_nb_epochs=1,
-        gpus=[0, 1],
-        distributed_backend='ddp',
-        use_amp=True
-    )
-
-    run_gpu_model_test(trainer_options, model, hparams)
-
 
 def test_cpu_restore_training():
     """
@@ -126,6 +99,36 @@ def test_cpu_restore_training():
     clear_save_dir()
 
 
+def test_amp_gpu_ddp():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+
+
+
 def test_cpu_slurm_save_load():
     """
     Verify model save/load/checkpoint on CPU

From a931ded310a20f3acde9a66dc4c9dd3b1f8cd48f Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:35:02 -0400
Subject: [PATCH 08/22] removed bad hook call

---
 pytorch_lightning/root_module/model_saving.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py
index 00b95460..a21f8a10 100644
--- a/pytorch_lightning/root_module/model_saving.py
+++ b/pytorch_lightning/root_module/model_saving.py
@@ -155,6 +155,8 @@ class TrainerIO(object):
         self.current_epoch = checkpoint['epoch']
 
         # restore the optimizers
+        import pdb
+        pdb.set_trace()
         optimizer_states = checkpoint['optimizer_states']
         for optimizer, opt_state in zip(self.optimizers, optimizer_states):
             optimizer.load_state_dict(opt_state)

From d5fd16a478148f6ec46ca73fc4f062f3563a14b9 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:38:45 -0400
Subject: [PATCH 09/22] removed bad hook call

---
 pytorch_lightning/models/trainer.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index 4d72618b..7aa1c037 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -244,9 +244,6 @@ class Trainer(TrainerIO):
             '''
             raise ModuleNotFoundError(msg)
 
-        # restore training and model
-        self.restore_state_if_existing_checkpoint()
-
     def restore_state_if_existing_checkpoint(self):
         # restore trainer state and model if there is a weight for this experiment
         last_epoch = -1
@@ -624,6 +621,9 @@ class Trainer(TrainerIO):
         ref_model.trainer = self
         ref_model.experiment = self.experiment
 
+        # restore training and model
+        self.restore_state_if_existing_checkpoint()
+
         # run tiny validation to make sure program won't crash during val
         _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
 

From 5c398d7a4ec9d2d1b6795f82d7e7355fc55230aa Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:39:41 -0400
Subject: [PATCH 10/22] removed bad hook call

---
 pytorch_lightning/root_module/model_saving.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py
index a21f8a10..00b95460 100644
--- a/pytorch_lightning/root_module/model_saving.py
+++ b/pytorch_lightning/root_module/model_saving.py
@@ -155,8 +155,6 @@ class TrainerIO(object):
         self.current_epoch = checkpoint['epoch']
 
         # restore the optimizers
-        import pdb
-        pdb.set_trace()
         optimizer_states = checkpoint['optimizer_states']
         for optimizer, opt_state in zip(self.optimizers, optimizer_states):
             optimizer.load_state_dict(opt_state)

From b0fae555718a14fac66b70b5f5e20e0b8ed5d65b Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:42:14 -0400
Subject: [PATCH 11/22] fixed restore location

---
 pytorch_lightning/models/trainer.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index 7aa1c037..cbec5228 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -621,9 +621,6 @@ class Trainer(TrainerIO):
         ref_model.trainer = self
         ref_model.experiment = self.experiment
 
-        # restore training and model
-        self.restore_state_if_existing_checkpoint()
-
         # run tiny validation to make sure program won't crash during val
         _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
 
@@ -635,8 +632,12 @@ class Trainer(TrainerIO):
         # if cluster resets state, the model will update with the saved weights
         self.model = model
 
+        # restore training and model before hpc call
+        self.restore_state_if_existing_checkpoint()
+
         # enable cluster checkpointing
         # also restores training state
+        # hpc checkpoint overrides any other checkpoints loaded before
         if self.cluster is not None:  # pragma: no cover
             self.enable_auto_hpc_walltime_manager()
 

From 8e4fe2002b50a3cb49976b38ad85ff09da89e622 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:45:57 -0400
Subject: [PATCH 12/22] fixed restore location

---
 tests/test_models.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/test_models.py b/tests/test_models.py
index 1708134a..6ee97ffe 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -75,6 +75,9 @@ def test_cpu_restore_training():
     new_exp = get_exp(False, version=test_exp_version)
     trainer_options = dict(
         max_nb_epochs=1,
+        val_check_interval=0.50,
+        val_percent_check=0.2,
+        train_percent_check=0.2,
         experiment=new_exp,
         checkpoint_callback=ModelCheckpoint(save_dir),
     )

From cdbcbad3529f1292339fe89950ba5c820f1f4727 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:51:55 -0400
Subject: [PATCH 13/22] added hook on_sanity_check_start

---
 pytorch_lightning/models/trainer.py    | 8 +++++---
 pytorch_lightning/root_module/hooks.py | 8 ++++++++
 tests/test_models.py                   | 2 +-
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index cbec5228..07dedd53 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -621,9 +621,6 @@ class Trainer(TrainerIO):
         ref_model.trainer = self
         ref_model.experiment = self.experiment
 
-        # run tiny validation to make sure program won't crash during val
-        _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
-
         # save exp to get started
         if self.proc_rank == 0:
             self.experiment.save()
@@ -641,9 +638,14 @@ class Trainer(TrainerIO):
         if self.cluster is not None:  # pragma: no cover
             self.enable_auto_hpc_walltime_manager()
 
+        # run tiny validation to make sure program won't crash during val
+        model.on_sanity_check_start()
+        _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
+
         # ---------------------------
         # CORE TRAINING LOOP
         # ---------------------------
+
         self.__train()
 
     def __train(self):
diff --git a/pytorch_lightning/root_module/hooks.py b/pytorch_lightning/root_module/hooks.py
index 849826a8..06ec614e 100644
--- a/pytorch_lightning/root_module/hooks.py
+++ b/pytorch_lightning/root_module/hooks.py
@@ -2,6 +2,14 @@ import torch
 
 
 class ModelHooks(torch.nn.Module):
+
+    def on_sanity_check_start(self):
+        """
+        Called before starting validate
+        :return:
+        """
+        pass
+
     def on_batch_start(self, data_batch):
         pass
 
diff --git a/tests/test_models.py b/tests/test_models.py
index 6ee97ffe..14f1e4ac 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -93,7 +93,7 @@ def test_cpu_restore_training():
         new_pred = trainer.model(x)
         assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
 
-    model.on_epoch_start = assert_pred_same
+    model.on_sanity_check_start = assert_pred_same
 
     # by calling fit again, we trigger training, loading weights from the cluster
     # and our hook to predict using current model before any more weight updates

From b8ce4adfa8988c503f2337e5aaa73d7e8c7e0281 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:55:47 -0400
Subject: [PATCH 14/22] debug

---
 pytorch_lightning/models/trainer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index 07dedd53..1ba85202 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -630,6 +630,8 @@ class Trainer(TrainerIO):
         self.model = model
 
         # restore training and model before hpc call
+        import pdb
+        pdb.set_trace()
         self.restore_state_if_existing_checkpoint()
 
         # enable cluster checkpointing

From 27e88fde31765872131a275255fcdd915f359615 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 07:57:03 -0400
Subject: [PATCH 15/22] debug

---
 pytorch_lightning/models/trainer.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index 1ba85202..07dedd53 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -630,8 +630,6 @@ class Trainer(TrainerIO):
         self.model = model
 
         # restore training and model before hpc call
-        import pdb
-        pdb.set_trace()
         self.restore_state_if_existing_checkpoint()
 
         # enable cluster checkpointing

From 1e17bf76aad4043eedc745fdb13ba9b13cea32c4 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 08:01:33 -0400
Subject: [PATCH 16/22] debug

---
 tests/test_models.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_models.py b/tests/test_models.py
index 14f1e4ac..9142ca45 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -42,7 +42,7 @@ def test_cpu_restore_training():
     exp.save()
 
     trainer_options = dict(
-        max_nb_epochs=1,
+        max_nb_epochs=2,
         val_check_interval=0.50,
         val_percent_check=0.2,
         train_percent_check=0.2,
@@ -53,7 +53,7 @@ def test_cpu_restore_training():
     # fit model
     trainer = Trainer(**trainer_options)
     result = trainer.fit(model)
-    real_global_step = trainer.global_step
+    real_global_epoch = trainer.current_epoch
 
     # traning complete
     assert result == 1, 'amp + ddp model failed to complete'
@@ -86,7 +86,7 @@ def test_cpu_restore_training():
 
     # set the epoch start hook so we can predict before the model does the full training
     def assert_pred_same():
-        assert trainer.global_step == real_global_step and trainer.global_step > 0
+        assert trainer.current_epoch == real_global_epoch and trainer.real_global_epoch > 0
 
         # predict with loaded model to make sure answers are the same
         trainer.model.eval()

From 2018380598b2c0916d198260be236c714b74e894 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 08:01:42 -0400
Subject: [PATCH 17/22] debug

---
 tests/test_models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_models.py b/tests/test_models.py
index 9142ca45..575ceb05 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -74,7 +74,7 @@ def test_cpu_restore_training():
     # we want to see if the weights come back correctly
     new_exp = get_exp(False, version=test_exp_version)
     trainer_options = dict(
-        max_nb_epochs=1,
+        max_nb_epochs=2,
         val_check_interval=0.50,
         val_percent_check=0.2,
         train_percent_check=0.2,

From 0527a1dad189278202de72f5afc0cf4c2fab1853 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 08:03:40 -0400
Subject: [PATCH 18/22] debug

---
 tests/test_models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_models.py b/tests/test_models.py
index 575ceb05..d3e694fc 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -86,7 +86,7 @@ def test_cpu_restore_training():
 
     # set the epoch start hook so we can predict before the model does the full training
     def assert_pred_same():
-        assert trainer.current_epoch == real_global_epoch and trainer.real_global_epoch > 0
+        assert trainer.current_epoch == real_global_epoch and trainer.current_epoch > 0
 
         # predict with loaded model to make sure answers are the same
         trainer.model.eval()

From 0b92fe6cea4d76a45eca97f517fefc0ee5114089 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 08:07:59 -0400
Subject: [PATCH 19/22] updated test

---
 tests/test_models.py | 21 +++++----------------
 1 file changed, 5 insertions(+), 16 deletions(-)

diff --git a/tests/test_models.py b/tests/test_models.py
index d3e694fc..d30a6795 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -58,17 +58,6 @@ def test_cpu_restore_training():
     # traning complete
     assert result == 1, 'amp + ddp model failed to complete'
 
-    # predict with trained model before saving
-    # make a prediction
-    for batch in model.test_dataloader:
-        break
-
-    x, y = batch
-    x = x.view(x.size(0), -1)
-
-    model.eval()
-    pred_before_saving = model(x)
-
     # wipe-out trainer and model
     # retrain with not much data... this simulates picking training back up after slurm
     # we want to see if the weights come back correctly
@@ -85,15 +74,15 @@ def test_cpu_restore_training():
     model = LightningTestModel(hparams)
 
     # set the epoch start hook so we can predict before the model does the full training
-    def assert_pred_same():
+    def assert_good_acc():
         assert trainer.current_epoch == real_global_epoch and trainer.current_epoch > 0
 
-        # predict with loaded model to make sure answers are the same
+        # if model and state loaded correctly, predictions will be good even though we
+        # haven't trained with the new loaded model
         trainer.model.eval()
-        new_pred = trainer.model(x)
-        assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
+        run_prediction(trainer.val_dataloader, trainer.model)
 
-    model.on_sanity_check_start = assert_pred_same
+    model.on_sanity_check_start = assert_good_acc
 
     # by calling fit again, we trigger training, loading weights from the cluster
     # and our hook to predict using current model before any more weight updates

From df5e10d0aa924881c07a77e461a2cee562d51bf2 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 08:12:54 -0400
Subject: [PATCH 20/22] updated test

---
 pytorch_lightning/models/trainer.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index 07dedd53..3f95d045 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -252,6 +252,10 @@ class Trainer(TrainerIO):
         # find last epoch
         checkpoints = os.listdir(self.checkpoint_callback.filepath)
         for name in checkpoints:
+            # ignore hpc ckpts
+            if 'hpc_' in name:
+                continue
+
             if '.ckpt' in name:
                 epoch = name.split('epoch_')[1]
                 epoch = int(re.sub('[^0-9]', '' ,epoch))

From d539e490856f27a298fb805a5859145a1e224f61 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 08:14:52 -0400
Subject: [PATCH 21/22] updated test

---
 pytorch_lightning/models/trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index 3f95d045..caf0f2a1 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -643,7 +643,7 @@ class Trainer(TrainerIO):
             self.enable_auto_hpc_walltime_manager()
 
         # run tiny validation to make sure program won't crash during val
-        model.on_sanity_check_start()
+        ref_model.on_sanity_check_start()
         _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
 
         # ---------------------------

From 2f1df17371f826a98964b0641427f9634bcef042 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 09:01:19 -0400
Subject: [PATCH 22/22] updated test

---
 pytorch_lightning/models/trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index caf0f2a1..94e4c7ab 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -1,5 +1,5 @@
 """
-The trainer handles all the logic for running a val loop, training loop, distributing, etc...
+The trainer handles all the logic for running a val loop, training loop, distributing, etc.. .
 """
 import subprocess
 import traceback