Repository: Lightning-AI/pytorch-lightning

Issue:
Bug in `LightningModule.load_from_checkpoint`
Suppose you have a class `Model(LightningModule)` (a `LightningModule` subclass) and as parameters it takes
`params: argparse.Namespace` (without a default value), along with other keyword arguments which might take default arguments. Loading from a checkpoint

```python
model = Model(params=params, **kwargs)
...  # saves a checkpoint
model = Model.load_from_checkpoint(params=params, **kwargs)  
# same behavior arises if params is passed as positional argument
```

will throw the error `TypeError: __init__() got multiple values for argument 'params'`.
This comes from the fact that in line `205` in `pytorch_lightning/core/saving.py`

```python
model = cls(*cls_args, **cls_kwargs)
```

when inspected, `cls_args` is `({'params': params},)` and `cls_kwargs` is, of course, `{'params': params}`.
This stems from the fact in lines `189-196` (in `saving.py`):

```python
            if args_name == 'kwargs':
                # in case the class cannot take any extra argument filter only the possible
                cls_kwargs.update(**model_args)
            elif args_name:
                if args_name in cls_init_args_name:
                    cls_kwargs.update({args_name: model_args})
            else:
                cls_args = (model_args,) + cls_args
```

the `else` clause is getting called and `model_args` contains a copy of kwargs!

To be honest, it's very unclear to me exactly what `model_args` is supposed to be here, but what is clear is that
this block needs more checking so that we don't just stick copies of items in `cls_kwargs` into `cls_args`. Just doing
something silly like replacing the `else` clause with

```python
from typing import Sequence
...
             elif isinstance(model_args, Sequence):
                cls_args = tuple(model_args) + cls_args
```

fixes the problem for me, but as I said, I'm unclear on exactly how model_args is being used, so I'm sure there is a better
solution.



" __init__() got multiple values for argument" when load_from_checkpoint
### What is your question?
I follow the docs [SAVING AND LOADING WEIGHTS](https://pytorch-lightning.readthedocs.io/en/stable/weights_loading.html#saving-and-loading-weights) but meet some problems.
When I load model from ckpt, I meet the error `__init__() got multiple values for argument 'net'`

### Code


Here is the [Colab version](https://colab.research.google.com/drive/13LAiSnTNXHl_BhethgFGljHfs2CCs4dF?usp=sharing) about the problem code.

#### code in colab
```python
import os
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms
import pytorch_lightning as pl
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self, h_dim):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, h_dim)
        self.l2 = torch.nn.Linear(h_dim, 10)

    def forward(self, x):
        x = torch.relu(self.l1(x.view(x.size(0), -1)))
        x = torch.relu(self.l2(x.view(x.size(0), -1)))
        return x

class MNISTModel(pl.LightningModule):

    def __init__(self, net, h_dim):
        super(MNISTModel, self).__init__()
        # self.save_hyperparameters()
        if net == 'cnn':
          self.net = CNN(h_dim)

    def forward(self, x):
        return self.net(x)

    def training_step(self, batch, batch_nb):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        tensorboard_logs = {'train_loss': loss}
        return {'loss': loss, 'log': tensorboard_logs}

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

train_loader = DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
mnist_model = MNISTModel('cnn', h_dim=200)
trainer = pl.Trainer(gpus=1, progress_bar_refresh_rate=20, max_epochs=3)    
trainer.fit(mnist_model, train_loader)  
trainer.save_checkpoint('mnist.ckpt')
```
when I want to load model from mnist.ckpt, I meet an error:
```python
model = MNISTModel.load_from_checkpoint('mnist.ckpt', net='cnn', h_dim=200)
```
```
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/core/saving.py in load_from_checkpoint(cls, checkpoint_path, map_location, hparams_file, tags_csv, *args, **kwargs)
    167         checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
    168 
--> 169         model = cls._load_model_state(checkpoint, *args, **kwargs)
    170         return model
    171 

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/core/saving.py in _load_model_state(cls, checkpoint, *cls_args, **cls_kwargs)
    203         if len(cls_spec.args) <= 1 and not cls_spec.kwonlyargs:
    204             cls_args, cls_kwargs = [], {}
--> 205         model = cls(*cls_args, **cls_kwargs)
    206         # load the state_dict on the model automatically
    207         model.load_state_dict(checkpoint['state_dict'])

TypeError: __init__() got multiple values for argument 'net'
```
### What's your environment?
pytorch-lightning: 0.8.4
Colab environment with GPU



Discussion:
Hi! thanks for your contribution!, great first issue!
Hi! thanks for your contribution!, great first issue!
I have checked it on different versions:
- 0.7.6: Ok
- 0.8.0: Ok
- 0.8.1 (and higher): TypeError

Could this behaviour have been introduced in 0.8.1?
same issue here! any workaround? i tried install 0.8.0 but it pip says it can't find that version
The problem is here:
https://github.com/PyTorchLightning/pytorch-lightning/blob/a34609ef0ef3c5c6fac5dd30520ba45dbe0c9159/pytorch_lightning/core/saving.py#L196
Since `args_name` is None so `model_args` will remain a dict. Changing it to `cls_kwargs.update(**model_args)` works.
@rohitgr7 can we work around this somehow?
@rohitgr7 try this: Lightning.load_from_checkpoint(path, model=model, kwargs=dict(param_name=param_value))
@GazizovMarat this should work but still should work simply by just passing the `args` or `kwargs` as normal params.
https://github.com/PyTorchLightning/pytorch-lightning/blob/7b4db3045dcc9e6bb0b66e409b25bb2c7fa378f0/pytorch_lightning/core/saving.py#L65
https://github.com/PyTorchLightning/pytorch-lightning/blob/7b4db3045dcc9e6bb0b66e409b25bb2c7fa378f0/pytorch_lightning/core/saving.py#L103
@rohitgr7  yes, it's just a workaround at the moment
I caught the same bug introduced in this issue.

[pytorch_lightning.core.saving.ModelIO.load_from_checkpoint()](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.core.saving.html#pytorch_lightning.core.saving.ModelIO.load_from_checkpoint) has `args` in the arguments to initialize `pl.LightningModule`.
If the module has multiple arguments, the method doesn't work correctly.

### Snippets not to add arguments for `pl.LightningModule`

```py
import pytorch_lightning as pl

class sample(pl.LightningModule):

    def __init__(self, one, two, three):
        super(sample, self).__init__()

    def forward(self, input, **kwargs):
        return 0

    def training_step(self, batch, batch_idx, optimizer_idx=0):
        return 0

    def validation_step(self, batch, batch_idx, optimizer_idx=0):
        return 0


if __name__ == "__main__":
    model = sample.load_from_checkpoint(checkpoint_path=path)  # please set `path` properly
    print(model)
```

Then, I got a following error.

```sh
Traceback (most recent call last):
  File ".\min.py", line 20, in <module>
    model = sample.load_from_checkpoint(checkpoint_path="...")
  File "...\lib\site-packages\pytorch_lightning\core\saving.py", line 169, in load_from_checkpoint
    model = cls._load_model_state(checkpoint, *args, **kwargs)
  File "...\lib\site-packages\pytorch_lightning\core\saving.py", line 205, in _load_model_state
    model = cls(*cls_args, **cls_kwargs)
TypeError: __init__() missing 2 required positional arguments: 'two' and 'three'
```

### Snippets to add arguments for `pl.LightningModule`

```py
import pytorch_lightning as pl

class sample(pl.LightningModule):

    def __init__(self, one, two, three):
        super(sample, self).__init__()

    def forward(self, input, **kwargs):
        return 0

    def training_step(self, batch, batch_idx, optimizer_idx=0):
        return 0

    def validation_step(self, batch, batch_idx, optimizer_idx=0):
        return 0


if __name__ == "__main__":
    model = sample.load_from_checkpoint(checkpoint_path=path, one=1, two=2, three=3) # please set `path` properly
    print(model)
```

Then, I got a following error.

```sh
Traceback (most recent call last):
  File ".\min.py", line 20, in <module>
    model = sample.load_from_checkpoint(checkpoint_path="...",
  File "...\lib\site-packages\pytorch_lightning\core\saving.py", line 169, in load_from_checkpoint
    model = cls._load_model_state(checkpoint, *args, **kwargs)
  File "...\lib\site-packages\pytorch_lightning\core\saving.py", line 205, in _load_model_state
    model = cls(*cls_args, **cls_kwargs)
TypeError: __init__() got multiple values for argument 'one'
```

### Environment

 - PyTorchLightning Version: 0.8.5
 - PyTorch Version: 1.5.1
 - OS: Windows
 - How you installed PyTorch: pip
 - Python version: 3.8.3
 - CUDA/cuDNN version: 10.2
 - GPU models and configuration: RTX 2080 Ti

Could you fix it?
Current workaround suggested here: https://github.com/PyTorchLightning/pytorch-lightning/issues/2550#issuecomment-659568926
Oh sorry, I missed the comments.
Thank you! I'll try it!
@Borda It seems there may be an issue here. In the [example above](https://github.com/PyTorchLightning/pytorch-lightning/issues/2550#issue-653041505), the `MNISTModel` is trained without `self.save_hyperparameters()` in `__init__`. When the user wants to load model weights from a checkpoint, the line `MNISTModel.load_from_checkpoint('mnist.ckpt', net='cnn', h_dim=200)` fails. As suggested by @rohitgr7, the issue seems to be in [`_load_model_state`](https://github.com/PyTorchLightning/pytorch-lightning/blob/1e16681693e113d25a4207ce2827a39c64b19211/pytorch_lightning/core/saving.py#L173). Indeed, in this example, at [line 196](https://github.com/PyTorchLightning/pytorch-lightning/blob/1e16681693e113d25a4207ce2827a39c64b19211/pytorch_lightning/core/saving.py#L196), `cls_args` equals `({'net': 'cnn', 'h_dim': 200},)` whereas `cls_kwargs` also equals `{'net': 'cnn', 'h_dim': 200}`. Later in this function, when `model = cls(*cls_args, **cls_kwargs)` is called, the class receives multiple values for the `net` and `h_dim` parameters. Hence, the error! 

The important point here is that one shall be able to load model weights from a checkpoint is the model was trained without `self.save_hyperparameters()`. As a fix for this issue, I propose to remove the following line in [`_load_model_state`](https://github.com/PyTorchLightning/pytorch-lightning/blob/1e16681693e113d25a4207ce2827a39c64b19211/pytorch_lightning/core/saving.py#L173)

```python
cls_args = (model_args,) + cls_args
```

as it seems to only duplicates the class arguments. I implemented this fix locally and all tests are passing ✅ ! 

Shall I make a PR ? 🚀 

cc @davidseroussi @invisprints
@jbschiratti please ping me in the PR if you make one :)

Resolving patch:
diff --git a/pytorch_lightning/core/saving.py b/pytorch_lightning/core/saving.py
--- a/pytorch_lightning/core/saving.py
+++ b/pytorch_lightning/core/saving.py
@@ -52,7 +52,6 @@ class ModelIO(object):
     def load_from_checkpoint(
         cls,
         checkpoint_path: str,
-        *args,
         map_location: Optional[Union[Dict[str, str], str, torch.device, int, Callable]] = None,
         hparams_file: Optional[str] = None,
         strict: bool = True,
@@ -66,7 +65,6 @@ def load_from_checkpoint(
 
         Args:
             checkpoint_path: Path to checkpoint. This can also be a URL.
-            args: Any positional args needed to init the model.
             map_location:
                 If your checkpoint saved a GPU model and you now load on CPUs
                 or a different number of GPUs, use this to map to the new setup.
@@ -152,44 +150,48 @@ def load_from_checkpoint(
         # override the hparams with values that were passed in
         checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
 
-        model = cls._load_model_state(checkpoint, *args, strict=strict, **kwargs)
+        model = cls._load_model_state(checkpoint, strict=strict, **kwargs)
         return model
 
     @classmethod
-    def _load_model_state(cls, checkpoint: Dict[str, Any], *cls_args, strict: bool = True, **cls_kwargs):
+    def _load_model_state(cls, checkpoint: Dict[str, Any], strict: bool = True, **cls_kwargs_new):
         cls_spec = inspect.getfullargspec(cls.__init__)
         cls_init_args_name = inspect.signature(cls.__init__).parameters.keys()
+
+        self_name = cls_spec.args[0]
+        drop_names = (self_name, cls_spec.varargs, cls_spec.varkw)
+        cls_init_args_name = list(filter(lambda n: n not in drop_names, cls_init_args_name))
+
+        cls_kwargs_loaded = {}
         # pass in the values we saved automatically
         if cls.CHECKPOINT_HYPER_PARAMS_KEY in checkpoint:
-            model_args = {}
 
-            # add some back compatibility, the actual one shall be last
-            for hparam_key in CHECKPOINT_PAST_HPARAMS_KEYS + (cls.CHECKPOINT_HYPER_PARAMS_KEY,):
-                if hparam_key in checkpoint:
-                    model_args.update(checkpoint[hparam_key])
+            # 1. (backward compatibility) Try to restore model hparams from checkpoint using old/past keys
+            for _old_hparam_key in CHECKPOINT_PAST_HPARAMS_KEYS:
+                cls_kwargs_loaded.update(checkpoint.get(_old_hparam_key, {}))
 
-            model_args = _convert_loaded_hparams(model_args, checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_TYPE))
+            # 2. Try to restore model hparams from checkpoint using the new key
+            _new_hparam_key = cls.CHECKPOINT_HYPER_PARAMS_KEY
+            cls_kwargs_loaded.update(checkpoint.get(_new_hparam_key))
 
+            # 3. Ensure that `cls_kwargs_old` has the right type, back compatibility between dict and Namespace
+            cls_kwargs_loaded = _convert_loaded_hparams(cls_kwargs_loaded, checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_TYPE))
+
+            # 4. Update cls_kwargs_new with cls_kwargs_old, such that new has higher priority
             args_name = checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_NAME)
+            if args_name and args_name in cls_init_args_name:
+                cls_kwargs_loaded = {args_name: cls_kwargs_loaded}
 
-            if args_name == 'kwargs':
-                # in case the class cannot take any extra argument filter only the possible
-                cls_kwargs.update(**model_args)
-            elif args_name:
-                if args_name in cls_init_args_name:
-                    cls_kwargs.update({args_name: model_args})
-            else:
-                cls_args = (model_args,) + cls_args
+        _cls_kwargs = {}
+        _cls_kwargs.update(cls_kwargs_loaded)
+        _cls_kwargs.update(cls_kwargs_new)
 
         if not cls_spec.varkw:
             # filter kwargs according to class init unless it allows any argument via kwargs
-            cls_kwargs = {k: v for k, v in cls_kwargs.items() if k in cls_init_args_name}
+            _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.varargs and not cls_spec.kwonlyargs:
-            cls_args, cls_kwargs = [], {}
+        model = cls(**_cls_kwargs)
 
-        model = cls(*cls_args, **cls_kwargs)
         # load the state_dict on the model automatically
         model.load_state_dict(checkpoint['state_dict'], strict=strict)
 
@@ -239,7 +241,7 @@ def on_hpc_load(self, checkpoint: Dict[str, Any]) -> None:
 
 
 def _convert_loaded_hparams(model_args: dict, hparams_type: Union[Callable, str] = None) -> object:
-    """Convert hparams according given type in callable or string (past) format"""
+    """Convert hparams according given type in callable or string (past) format."""
     # if not hparams type define
     if not hparams_type:
         return model_args


Test patch:
diff --git a/tests/base/model_template.py b/tests/base/model_template.py
--- a/tests/base/model_template.py
+++ b/tests/base/model_template.py
@@ -39,17 +39,17 @@ class EvalModelTemplate(
     """
 
     def __init__(
-        self,
-        drop_prob: float = 0.2,
-        batch_size: int = 32,
-        in_features: int = 28 * 28,
-        learning_rate: float = 0.001 * 8,
-        optimizer_name: str = 'adam',
-        data_root: str = PATH_DATASETS,
-        out_features: int = 10,
-        hidden_dim: int = 1000,
-        b1: float = 0.5,
-        b2: float = 0.999,
+            self,
+            drop_prob: float = 0.2,
+            batch_size: int = 32,
+            in_features: int = 28 * 28,
+            learning_rate: float = 0.001 * 8,
+            optimizer_name: str = 'adam',
+            data_root: str = PATH_DATASETS,
+            out_features: int = 10,
+            hidden_dim: int = 1000,
+            b1: float = 0.5,
+            b2: float = 0.999,
     ):
         # init superclass
         super().__init__()
diff --git a/tests/models/test_hparams.py b/tests/models/test_hparams.py
--- a/tests/models/test_hparams.py
+++ b/tests/models/test_hparams.py
@@ -538,5 +538,6 @@ def test_args(tmpdir):
     trainer.fit(model)
 
     raw_checkpoint_path = _raw_checkpoint_path(trainer)
-    model = SubClassVarArgs.load_from_checkpoint(raw_checkpoint_path)
-    assert model.hparams == hparams
+    with pytest.raises(TypeError, match="__init__\(\) got an unexpected keyword argument 'test'"):
+        SubClassVarArgs.load_from_checkpoint(raw_checkpoint_path)
+
diff --git a/tests/models/test_restore.py b/tests/models/test_restore.py
--- a/tests/models/test_restore.py
+++ b/tests/models/test_restore.py
@@ -2,17 +2,18 @@
 import logging as log
 import os
 import pickle
-import functools
 
 import cloudpickle
 import pytest
 import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
 
 import tests.base.develop_pipelines as tpipes
 import tests.base.develop_utils as tutils
-from pytorch_lightning import Trainer, Callback
+from pytorch_lightning import Trainer, LightningModule, Callback
 from pytorch_lightning.callbacks import ModelCheckpoint
-from tests.base import EvalModelTemplate, GenericEvalModelTemplate
+from tests.base import EvalModelTemplate, GenericEvalModelTemplate, TrialMNIST
 
 
 class ModelTrainerPropertyParity(Callback):
@@ -217,6 +218,12 @@ def test_load_model_from_checkpoint(tmpdir, model_template):
 
     # load last checkpoint
     last_checkpoint = sorted(glob.glob(os.path.join(trainer.checkpoint_callback.dirpath, "*.ckpt")))[-1]
+
+    # Since `EvalModelTemplate` has `_save_hparams = True` by default, check that ckpt has hparams
+    ckpt = torch.load(last_checkpoint)
+    assert model_template.CHECKPOINT_HYPER_PARAMS_KEY in ckpt.keys(), 'module_arguments missing from checkpoints'
+
+    # Ensure that model can be correctly restored from checkpoint
     pretrained_model = model_template.load_from_checkpoint(last_checkpoint)
 
     # test that hparams loaded correctly
@@ -227,6 +234,7 @@ def test_load_model_from_checkpoint(tmpdir, model_template):
     for (old_name, old_p), (new_name, new_p) in zip(model.named_parameters(), pretrained_model.named_parameters()):
         assert torch.all(torch.eq(old_p, new_p)), 'loaded weights are not the same as the saved weights'
 
+    # Check `test` on pretrained model:
     new_trainer = Trainer(**trainer_options)
     new_trainer.test(pretrained_model)