Repository: Lightning-AI/pytorch-lightning

Issue:
Save/load hyperparameters with checkpoint
Right now, test tube users can have `model.load_from_metrics`, but there's no equivalent functionality for users of other logging frameworks. I believe we discussed a while ago that this functionality should be decoupled from logggers.

Proposal:
* When hyperparameters are stored as an `argparse.Namespace` in the `hparams` property of the model (as already required to get automatic hyperparam logging), save these parameters as part of checkpoints.
* Add a `load_from_checkpoint` classmethod to LightningModule that basically duplicates the functionality of `load_from_metrics`, but gets the hyperparameters from the checkpoint directly.

I'm happy to send a PR if this proposal looks okay.


Discussion:
yeah that makes sense.

So in the checkpoint do:

```python
checkpoint['hparams'] = ...
```
?

That should be the standard going forward. Logging is nice for Testtube UI purposes to see what went into an experiment. But that's separate from loading weights imho.

Resolving patch:
diff --git a/pl_examples/basic_examples/lightning_module_template.py b/pl_examples/basic_examples/lightning_module_template.py
--- a/pl_examples/basic_examples/lightning_module_template.py
+++ b/pl_examples/basic_examples/lightning_module_template.py
@@ -158,7 +158,7 @@ def validation_end(self, outputs):
             val_loss = output['val_loss']
 
             # reduce manually when using dp
-            if self.trainer.use_dp:
+            if self.trainer.use_dp or self.trainer.use_ddp2:
                 val_loss = torch.mean(val_loss)
             val_loss_mean += val_loss
 
diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py
--- a/pytorch_lightning/root_module/root_module.py
+++ b/pytorch_lightning/root_module/root_module.py
@@ -1,4 +1,5 @@
 import warnings
+from argparse import Namespace
 
 import torch
 
@@ -177,6 +178,36 @@ def load_from_metrics(cls, weights_path, tags_csv):
 
         return model
 
+    @classmethod
+    def load_from_checkpoint(cls, checkpoint_path):
+        """
+        Primary way of loading model from a checkpoint
+        :param checkpoint_path:
+        :param map_location: dic for mapping storage {'cuda:1':'cuda:0'}
+        :return:
+        """
+
+        # load on CPU only to avoid OOM issues
+        # then its up to user to put back on GPUs
+        checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
+        try:
+            ckpt_hparams = checkpoint['hparams']
+        except KeyError:
+            raise IOError(
+                "Checkpoint does not contain hyperparameters. Are your model hyperparameters stored"
+                "in self.hparams?"
+            )
+        hparams = Namespace(**ckpt_hparams)
+
+        # load the state_dict on the model automatically
+        model = cls(hparams)
+        model.load_state_dict(checkpoint['state_dict'])
+
+        # give model a chance to load something
+        model.on_load_checkpoint(checkpoint)
+
+        return model
+
     def summarize(self, mode):
         model_summary = ModelSummary(self, mode=mode)
         print(model_summary)
diff --git a/pytorch_lightning/trainer/trainer_io.py b/pytorch_lightning/trainer/trainer_io.py
--- a/pytorch_lightning/trainer/trainer_io.py
+++ b/pytorch_lightning/trainer/trainer_io.py
@@ -1,6 +1,7 @@
 import os
 import re
 import signal
+import warnings
 from subprocess import call
 
 import torch
@@ -172,9 +173,16 @@ def dump_checkpoint(self):
 
         checkpoint['lr_schedulers'] = lr_schedulers
 
-        # add the state_dict from the model
+        # add the hparams and state_dict from the model
         model = self.get_model()
         checkpoint['state_dict'] = model.state_dict()
+        if hasattr(model, "hparams"):
+            checkpoint['hparams'] = vars(model.hparams)
+        else:
+            warnings.warn(
+                "Did not find hyperparameters at model.hparams. Saving checkpoint without"
+                " hyperparameters"
+            )
 
         # give the model a chance to add a few things
         model.on_save_checkpoint(checkpoint)


Test patch:
diff --git a/.run_local_tests.sh b/.run_local_tests.sh
--- a/.run_local_tests.sh
+++ b/.run_local_tests.sh
@@ -3,5 +3,6 @@ rm -rf _ckpt_*
 rm -rf tests/save_dir*
 rm -rf tests/mlruns_*
 rm -rf tests/tests/*
+rm -rf lightning_logs
 coverage run --source pytorch_lightning -m py.test pytorch_lightning tests pl_examples -v --doctest-modules
 coverage report -m
diff --git a/pytorch_lightning/testing/lm_test_module_mixins.py b/pytorch_lightning/testing/lm_test_module_mixins.py
--- a/pytorch_lightning/testing/lm_test_module_mixins.py
+++ b/pytorch_lightning/testing/lm_test_module_mixins.py
@@ -80,13 +80,13 @@ def validation_end(self, outputs):
             val_loss = output['val_loss']
 
             # reduce manually when using dp
-            if self.trainer.use_dp:
+            if self.trainer.use_dp or self.trainer.use_ddp2:
                 val_loss = torch.mean(val_loss)
             val_loss_mean += val_loss
 
             # reduce manually when using dp
             val_acc = output['val_acc']
-            if self.trainer.use_dp:
+            if self.trainer.use_dp or self.trainer.use_ddp2:
                 val_acc = torch.mean(val_acc)
 
             val_acc_mean += val_acc
diff --git a/tests/test_models.py b/tests/test_models.py
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -402,6 +402,47 @@ def test_running_test_pretrained_model():
     clear_save_dir()
 
 
+def test_load_model_from_checkpoint():
+    reset_seed()
+
+    """Verify test() on pretrained model"""
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        checkpoint_callback=True,
+        logger=False,
+        default_save_path=save_dir
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'training failed to complete'
+    pretrained_model = LightningTestModel.load_from_checkpoint(
+        os.path.join(trainer.checkpoint_callback.filepath, "_ckpt_epoch_1.ckpt")
+    )
+
+    # test that hparams loaded correctly
+    for k, v in vars(hparams).items():
+        assert getattr(pretrained_model.hparams, k) == v
+
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    # test we have good test accuracy
+    assert_ok_test_acc(new_trainer)
+    clear_save_dir()
+
+
 def test_running_test_pretrained_model_dp():
     reset_seed()