Repository: Lightning-AI/pytorch-lightning

Issue:
Streamlined UX in saving, loading, continue training.
Currently each time I ran the command with the same experiment name, a new version is created and trained from scratch.
If I wanted to load back the model to continue training, I have to create a different script, input the path for the previous checkpoints and meta_tags.

Proposal:
- Start training, input only 1 `log_path` and 1 `experiment_name`. Test_tube data and saved models will be included in the same folder.
- In each experiment, default is to create a new model and trained from scratch (current behavior)
- If user passed in `--continue-training`, model will load the latest model (not necessarily the best model) from the latest version, then create a new experiment version and continue training from there.
- If user passed in `--continue-training--best`, model will load the the best model from the latest version, then create a new experiment version and continue training from there.
- If user passed in `--continue-training(--best) --version=X`, auto load the model from version X and start training.

In the end, all the path and dir are determined from only the log_path and name of experiment, and default behavior is the same as of now: 
`python train.py --exp_name=exp1 --....`
If user want to change some hyperparams, or start finetuning, they will only need to add:
`python train.py --exp_name=exp1 --....  --continue-trainig--best --version=x`

What do you think?
Happy to discuss more, as I am not sure this is suitable for cluster training or not.
I can make a PR if you're interest.


Discussion:
Great suggestion. PR for this would be great, but let's do the feature in a slightly different way.

Let's pick 1 way of doing it with these considerations.

1. Test-tube logs and weights can't be forced to be in the same path. A user might have a restriction to save weights and logs to different folder. This could be for disk space or for ease of SCP files (especially true on corporate or .edu clusters)

2. With that said, continue training can still work. It's actually already supported but meant for cluster training only. However, I agree it's a good idea to change the signatures and docs a bit to ALSO allow this for non-cluster training.

To do that here are the things that need to be changed:
1. Update the [save function](https://github.com/williamFalcon/pytorch-lightning/blob/b88307e927e5240f711ab199dc81054e65965f67/pytorch_lightning/root_module/model_saving.py#L54) to behave like this [hpc_save function](https://github.com/williamFalcon/pytorch-lightning/blob/b88307e927e5240f711ab199dc81054e65965f67/pytorch_lightning/root_module/model_saving.py#L150-L175) (which saves training state as well).   
2. Delete the hpc_save function.
3. Register the new save fx [here](https://github.com/williamFalcon/pytorch-lightning/blob/b88307e927e5240f711ab199dc81054e65965f67/pytorch_lightning/root_module/model_saving.py#L106) (for cluster tng).
4. Change the signature of [hpc_load](https://github.com/williamFalcon/pytorch-lightning/blob/b88307e927e5240f711ab199dc81054e65965f67/pytorch_lightning/root_module/model_saving.py#L177).   
5. Register the new hpc_save in the same fx as (3).   
6. This is where it gets tricky. It might be easier to give the trainer either a new flag or a function like restore_training. Not sure which one might be a better UX. 

Let's pick one of these two:

A)
```python
trainer = Trainer(...)
trainer.restore_training(weights_path)
```

B)    
Or if the version is the same in the exp, it'll just pick right back up:
```python
exp = Experiment(version=same_as_before)
trainer = Trainer(experiment=exp)
```

7. Change the relevant tests.   
8. Update docs.   

At the end of the day, we want hpc_save/load and non-hpc_save/load to now behave the same.

Thoughts?

@cinjon any suggestions?
Great. I will start reading the hpc code as I didn't use DDP before and try making some changes.
ok, so, let’s do it where if the test tube version is the same then it continues training?

so, nothing of the external api needs to change, just internal + docs
@lkhphuc how's it going? Did you take a look at the hpc method signatures?
Sorry, I've been busy the last few days. I will spend some time on it in the next week.
Around the API for this change, I think:
- Everytime you run an experiment, the version of that experiment is bumped (current behavior).
- If additional argument `--load-model=best/latest`, the best/latest checkpoint from the *latest* version is loaded back.
- If additional arguments `--load-model=best/latest --from-version=X`, the best/latest checkpoint from version *X* is loaded back.
- If additional argument `--from-version=X`: the *latest* checkpoint from version *X* is loaded.

Together, I will introduce two new arguments in Trainer, the default will behave like current, combine they can load checkpoint best/latest  from version latest/X.

How do you think?
Should `--load-model=best/latest/X` with X is arbitrary epoch number?
No worries. I think we're saying the same thing but in different ways haha.

I don't think we need to add any more arguments. The only change to the public API that needs to be made is to load the latest weights whenever the exp version is specified.

Case 1: No exp version given. Version is bumped (current behavior)

Case 2: Exp version is given. Then exp continues where it left off (this is already the current behavior). The only addition here is  that it should load the latest weights available for that version.

Case 2 covers the 3 flags you mentioned, so there's no need to add any more arguments, just change the internal loading (ie: look for the latest checkpoint at the beginning of training. If the version is new there won't be any, if it's an old version, the user had to have set that on purpose. and thus would be expected to load the latest weights).


Resolving patch:
diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
--- 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
@@ -244,6 +244,32 @@ def __init__(self,
             '''
             raise ModuleNotFoundError(msg)
 
+    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:
+            # ignore hpc ckpts
+            if 'hpc_' in name:
+                continue
+
+            if '.ckpt' in name:
+                epoch = name.split('epoch_')[1]
+                epoch = int(re.sub('[^0-9]', '' ,epoch))
+
+                if epoch > last_epoch:
+                    last_epoch = epoch
+                    last_ckpt_name = name
+
+        # restore last checkpoint
+        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):
         return self.use_dp or self.use_ddp
@@ -599,9 +625,6 @@ def __run_pretrain_routine(self, model):
         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()
@@ -610,14 +633,23 @@ def __run_pretrain_routine(self, model):
         # 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()
 
+        # run tiny validation to make sure program won't crash during val
+        ref_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
--- a/pytorch_lightning/root_module/hooks.py
+++ b/pytorch_lightning/root_module/hooks.py
@@ -2,6 +2,14 @@
 
 
 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/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py
--- a/pytorch_lightning/root_module/model_saving.py
+++ b/pytorch_lightning/root_module/model_saving.py
@@ -58,6 +58,22 @@ def save_checkpoint(self, filepath):
         # 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'])
+
     def dump_checkpoint(self):
 
         checkpoint = {
@@ -198,15 +214,15 @@ def hpc_load(self, folderpath, on_gpu):
         # 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))
 


Test patch:
diff --git a/tests/test_models.py b/tests/test_models.py
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -24,6 +24,73 @@
 # ------------------------------------------------------------------------
 # TESTS
 # ------------------------------------------------------------------------
+
+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=2,
+        val_check_interval=0.50,
+        val_percent_check=0.2,
+        train_percent_check=0.2,
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    real_global_epoch = trainer.current_epoch
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # 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=2,
+        val_check_interval=0.50,
+        val_percent_check=0.2,
+        train_percent_check=0.2,
+        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_good_acc():
+        assert trainer.current_epoch == real_global_epoch and 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
+        trainer.model.eval()
+        run_prediction(trainer.val_dataloader, trainer.model)
+
+    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
+    trainer.fit(model)
+
+    clear_save_dir()
+
+
 def test_amp_gpu_ddp():
     """
     Make sure DDP + AMP work
@@ -52,6 +119,8 @@ def test_amp_gpu_ddp():
     run_gpu_model_test(trainer_options, model, hparams)
 
 
+
+
 def test_cpu_slurm_save_load():
     """
     Verify model save/load/checkpoint on CPU
@@ -610,10 +679,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