Fix hparams loading for model that accepts *args (Lightning-AI#2911)

* fix hparams loading for model that accepts *args

* add test case

* changelog

* pep

* fix test

Co-authored-by: William Falcon <[email protected]>

Author: 2 people

CHANGELOG.md

@@ -123,6 +123,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed save apex scaler states ([#2828](https://github.com/PyTorchLightning/pytorch-lightning/pull/2828))
 
+- Fixed a model loading issue with inheritance and variable positional arguments ([#2911](https://github.com/PyTorchLightning/pytorch-lightning/pull/2911))
+
 - Fixed passing `non_blocking=True` when transferring a batch object that does not support it ([#2910](https://github.com/PyTorchLightning/pytorch-lightning/pull/2910))
 
 ## [0.8.5] - 2020-07-09

pytorch_lightning/core/saving.py

@@ -167,8 +167,9 @@ def _load_model_state(cls, checkpoint: Dict[str, Any], *cls_args, **cls_kwargs):
             cls_kwargs = {k: v for k, v in cls_kwargs.items() if k in cls_init_args_name}
 
         # prevent passing positional arguments if class does not accept any
-        if len(cls_spec.args) <= 1 and not cls_spec.kwonlyargs:
+        if len(cls_spec.args) <= 1 and not cls_spec.varargs and not cls_spec.kwonlyargs:
             cls_args, cls_kwargs = [], {}
+
         model = cls(*cls_args, **cls_kwargs)
         # load the state_dict on the model automatically
         model.load_state_dict(checkpoint['state_dict'])

tests/models/test_hparams.py

@@ -497,3 +497,29 @@ def __init__(self, batch_size=15):
     raw_checkpoint_path = _raw_checkpoint_path(trainer)
     model = LocalModel.load_from_checkpoint(raw_checkpoint_path, non_exist_kwarg=99)
     assert 'non_exist_kwarg' not in model.hparams
+
+
+class SuperClassPositionalArgs(EvalModelTemplate):
+
+    def __init__(self, hparams):
+        super().__init__()
+        self._hparams = None  # pretend EvalModelTemplate did not call self.save_hyperparameters()
+        self.hparams = hparams
+
+
+class SubClassVarArgs(SuperClassPositionalArgs):
+    """ Loading this model should accept hparams and init in the super class """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+
+def test_args(tmpdir):
+    """ Test for inheritance: super class takes positional arg, subclass takes varargs. """
+    hparams = dict(test=1)
+    model = SubClassVarArgs(hparams)
+    trainer = Trainer(default_root_dir=tmpdir, max_epochs=1)
+    trainer.fit(model)
+
+    raw_checkpoint_path = _raw_checkpoint_path(trainer)
+    model = SubClassVarArgs.load_from_checkpoint(raw_checkpoint_path)
+    assert model.hparams == hparams