Repository: Lightning-AI/pytorch-lightning

Issue:
Error when loading with LitModel.load_from_checkpoint(path)?
## ❓ Questions and Help

### Before asking:   
1. search the issues.   Done
2. search the docs.    Done

<!-- If you still can't find what you need: -->

#### What is your question?

When I'm loading a Lightning model with LitModel.load_from_checkpoint(path) I always get this error:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-11-5cfc485836df> in <module>
      1 # OPTION 2:
      2 # test after loading weights
----> 3 model = LitModel.load_from_checkpoint("lightning_logs/version_0/checkpoints/epoch=2.ckpt")

~/anaconda3/envs/pytorch-lightning/lib/python3.7/site-packages/pytorch_lightning/core/saving.py in load_from_checkpoint(cls, checkpoint_path, map_location, hparams_file, tags_csv, *args, **kwargs)
    169         checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
    170 
--> 171         model = cls._load_model_state(checkpoint, *args, **kwargs)
    172         return model
    173 

~/anaconda3/envs/pytorch-lightning/lib/python3.7/site-packages/pytorch_lightning/core/saving.py in _load_model_state(cls, checkpoint, *args, **kwargs)
    195 
    196         # load the state_dict on the model automatically
--> 197         model = cls(*args, **kwargs)
    198         model.load_state_dict(checkpoint['state_dict'])
    199 

TypeError: __init__() takes 1 positional argument but 2 were given
```

#### Code

Simple example is from https://pytorch-lightning.readthedocs.io/en/latest/new-project.html.
Custom LightModul is very similar but I am using other data, sets, and loaders and a torchvision ResNet.


#### What have you tried?

I tried different checkpoints as well as relative and absolute paths.
Checkpoints are loaded from `lightning_logs/version_0/checkpoints/` which where saved automatically after stopping training.

Setup and training itself works without problems.

I also get the same error with the more complex setup.

Maybe I don't see something super obvious?

Thank you very much & kind regards


#### What's your environment?

 - OS: Ubuntu 18.04
 - Packaging: Conda
 - Version:
pytorch_select 0.2
pytorch 1.3.1
pytorch-lightning 0.8.0

An Extra argument passed to the class, loaded from load_from_checkpoint.
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

Hello,
I was facing few issues while using the `trainer.test()` function, on debugging I found out that the problem was with the `_load_model_state` class method which is called by `load_from_checkpoint`.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code For reference
<!-- Ideally attach a minimal code sample to reproduce the decried issue. 
Minimal means having the shortest code but still preserving the bug. -->
```
@classmethod
def _load_model_state(cls, checkpoint: Dict[str, Any], *args, **kwargs):
    # 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])

        if cls.CHECKPOINT_HYPER_PARAMS_TYPE in checkpoint:
            model_args = checkpoint[cls.CHECKPOINT_HYPER_PARAMS_TYPE](model_args)

        args_name = checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_NAME)
        init_args_name = inspect.signature(cls).parameters.keys()

        if args_name == 'kwargs':
            cls_kwargs = {k: v for k, v in model_args.items() if k in init_args_name}
            kwargs.update(**cls_kwargs)
        elif args_name:
            if args_name in init_args_name:
                kwargs.update({args_name: model_args})
        else:
            args = (model_args, ) + args

    # load the state_dict on the model automatically
    model = cls(*args, **kwargs)
    model.load_state_dict(checkpoint['state_dict'])

    # give model a chance to load something
    model.on_load_checkpoint(checkpoint)

    return model
```
Consider the case where the  model has no arguments, which corresponds to `LightModel.load_from_checkpoint('path')`. Here, the else clause of the if-elif is being executed where the `agrs` variable is updated from an empty tuple to a tuple with an empty dictionary `args = (model_args, ) + args` (as `model_args={}`). Therefore, while unpacking the args and kwargs (`model = cls(*args, **kwargs)`), There is an extra argument being passed which raises a `TypeError: __init__() takes 1 positional arguments but 2 were given`. #2364 

In some cases if the model has an argument and the user has forgotten to add it in the load_from_checkpoint, then an empty dictionary will be passed instead and it raises other errors depending on the code. For example, in the issue #2359 an empty dict is passed while loading the model and hence raises `RuntimeError: Error(s) in loading state_dict for Model:`.

I do not fully understand what is happening in the function. It would be great if someone can suggest changes to make in the comments so that I can start working after updating the changes in my forked repo.

#### Steps to reproduce
```
!pip install git+https://github.com/PytorchLightning/pytorch-lightning.git@master --upgrade

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

class MNISTModel(pl.LightningModule):

    def __init__(self):
        super(MNISTModel, self).__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    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 test_step(self, batch, batch_nb):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, 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()
trainer = pl.Trainer(gpus=1,max_epochs=3)    
trainer.fit(mnist_model, train_loader)  

test_loader = DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32)
trainer.test(test_dataloaders=test_loader)
```
Which returns:
```
TypeError                                 Traceback (most recent call last)

<ipython-input-5-50449ee4f6cc> in <module>()
      1 test_loader = DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32)
----> 2 trainer.test(test_dataloaders=test_loader)

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py in test(self, model, test_dataloaders, ckpt_path)
   1168             if ckpt_path == 'best':
   1169                 ckpt_path = self.checkpoint_callback.best_model_path
-> 1170             model = self.get_model().load_from_checkpoint(ckpt_path)
   1171 
   1172         self.testing = True

/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)
    201 
    202         # load the state_dict on the model automatically
--> 203         model = cls(*cls_args, **cls_kwargs)
    204         model.load_state_dict(checkpoint['state_dict'])
    205 

TypeError: __init__() takes 1 positional argument but 2 were given
```
#### Expected behavior
Start testing


Discussion:
Hi! thanks for your contribution!, great first issue!
Might be related to #2334 
If you put an hparams argument in your LitModel does it work? What about using 0.8.1?
Thank you for your reply!

This code snippet gives me the same error:
```
model = LitModel.load_from_checkpoint(checkpoint_path="lightning_logs/version_0/checkpoints/epoch=2.ckpt",
                                      hparams_file="lightning_logs/version_0/hparams.yaml")
```
Via conda I get no update and also tried to upgrade with pip but still 0.8.
Do I need for 0.8.1 an pip editable install via the GitHub repo?
Sorry, I meant changing the __init__ to (self, hparams).

To me it looks the load_from_checkpoint is trying to force hparams (2 arguments given)
But this might have been fixed in 0.8.1.

Try to `pip uninstall pytorch-lightning` and/or `pip install pytorch-lightning --upgrade`
Thank you, I was able to install the 0.8.1 version with this trick.

When I add `hparams` to the `def __init__(self, hparams)` to my lightning module I have to supply it during object creation too?

When I load the checkpoint with `torch.load` everything seems to be there.

I suspect the `args` supplied to `cls._load_model_state(checkpoint, *args, **kwargs)` are the reason this happens (`kwargs` are empty), will check later in detail.
Hi! thanks for your contribution!, great first issue!

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
@@ -171,6 +171,8 @@ def load_from_checkpoint(
 
     @classmethod
     def _load_model_state(cls, checkpoint: Dict[str, Any], *cls_args, **cls_kwargs):
+        cls_spec = inspect.getfullargspec(cls.__init__)
+        cls_init_args_name = inspect.signature(cls).parameters.keys()
         # pass in the values we saved automatically
         if cls.CHECKPOINT_HYPER_PARAMS_KEY in checkpoint:
             model_args = {}
@@ -183,14 +185,9 @@ def _load_model_state(cls, checkpoint: Dict[str, Any], *cls_args, **cls_kwargs):
             model_args = _convert_loaded_hparams(model_args, checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_TYPE))
 
             args_name = checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_NAME)
-            cls_spec = inspect.getfullargspec(cls.__init__)
-            kwargs_identifier = cls_spec.varkw
-            cls_init_args_name = inspect.signature(cls).parameters.keys()
 
             if args_name == 'kwargs':
                 # in case the class cannot take any extra argument filter only the possible
-                if not kwargs_identifier:
-                    model_args = {k: v for k, v in model_args.items() if k in cls_init_args_name}
                 cls_kwargs.update(**model_args)
             elif args_name:
                 if args_name in cls_init_args_name:
@@ -198,8 +195,15 @@ def _load_model_state(cls, checkpoint: Dict[str, Any], *cls_args, **cls_kwargs):
             else:
                 cls_args = (model_args,) + cls_args
 
-        # load the state_dict on the model automatically
+        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}
+
+        # prevent passing positional arguments if class does not accept any
+        if len(cls_spec.args) <= 1 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'])
 
         # give model a chance to load something


Test patch:
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
@@ -6,11 +6,13 @@
 import pytest
 import torch
 from omegaconf import OmegaConf, Container
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
 
 from pytorch_lightning import Trainer, LightningModule
 from pytorch_lightning.core.saving import save_hparams_to_yaml, load_hparams_from_yaml
 from pytorch_lightning.utilities import AttributeDict
-from tests.base import EvalModelTemplate
+from tests.base import EvalModelTemplate, TrialMNIST
 
 
 class SaveHparamsModel(EvalModelTemplate):
@@ -103,16 +105,16 @@ def test_explicit_args_hparams(tmpdir):
     """
 
     # define model
-    class TestModel(EvalModelTemplate):
+    class LocalModel(EvalModelTemplate):
         def __init__(self, test_arg, test_arg2):
             super().__init__()
             self.save_hyperparameters('test_arg', 'test_arg2')
 
-    model = TestModel(test_arg=14, test_arg2=90)
+    model = LocalModel(test_arg=14, test_arg2=90)
 
     # run standard test suite
-    raw_checkpoint_path = _run_standard_hparams_test(tmpdir, model, TestModel)
-    model = TestModel.load_from_checkpoint(raw_checkpoint_path, test_arg2=120)
+    raw_checkpoint_path = _run_standard_hparams_test(tmpdir, model, LocalModel)
+    model = LocalModel.load_from_checkpoint(raw_checkpoint_path, test_arg2=120)
 
     # config specific tests
     assert model.hparams.test_arg2 == 120
@@ -124,16 +126,16 @@ def test_implicit_args_hparams(tmpdir):
     """
 
     # define model
-    class TestModel(EvalModelTemplate):
+    class LocalModel(EvalModelTemplate):
         def __init__(self, test_arg, test_arg2):
             super().__init__()
             self.save_hyperparameters()
 
-    model = TestModel(test_arg=14, test_arg2=90)
+    model = LocalModel(test_arg=14, test_arg2=90)
 
     # run standard test suite
-    raw_checkpoint_path = _run_standard_hparams_test(tmpdir, model, TestModel)
-    model = TestModel.load_from_checkpoint(raw_checkpoint_path, test_arg2=120)
+    raw_checkpoint_path = _run_standard_hparams_test(tmpdir, model, LocalModel)
+    model = LocalModel.load_from_checkpoint(raw_checkpoint_path, test_arg2=120)
 
     # config specific tests
     assert model.hparams.test_arg2 == 120
@@ -145,12 +147,12 @@ def test_explicit_missing_args_hparams(tmpdir):
     """
 
     # define model
-    class TestModel(EvalModelTemplate):
+    class LocalModel(EvalModelTemplate):
         def __init__(self, test_arg, test_arg2):
             super().__init__()
             self.save_hyperparameters('test_arg')
 
-    model = TestModel(test_arg=14, test_arg2=90)
+    model = LocalModel(test_arg=14, test_arg2=90)
 
     # test proper property assignments
     assert model.hparams.test_arg == 14
@@ -166,7 +168,7 @@ def __init__(self, test_arg, test_arg2):
     assert raw_checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_KEY]['test_arg'] == 14
 
     # verify that model loads correctly
-    model = TestModel.load_from_checkpoint(raw_checkpoint_path, test_arg2=123)
+    model = LocalModel.load_from_checkpoint(raw_checkpoint_path, test_arg2=123)
     assert model.hparams.test_arg == 14
     assert 'test_arg2' not in model.hparams  # test_arg2 is not registered in class init
 
@@ -427,3 +429,71 @@ def test_hparams_save_yaml(tmpdir):
 
     save_hparams_to_yaml(path_yaml, OmegaConf.create(hparams))
     assert load_hparams_from_yaml(path_yaml) == hparams
+
+
+class NoArgsSubClassEvalModel(EvalModelTemplate):
+    def __init__(self):
+        super().__init__()
+
+
+class SimpleNoArgsModel(LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.l1 = torch.nn.Linear(28 * 28, 10)
+
+    def forward(self, x):
+        return torch.relu(self.l1(x.view(x.size(0), -1)))
+
+    def training_step(self, batch, batch_nb):
+        x, y = batch
+        loss = F.cross_entropy(self(x), y)
+        return {'loss': loss, 'log': {'train_loss': loss}}
+
+    def test_step(self, batch, batch_nb):
+        x, y = batch
+        loss = F.cross_entropy(self(x), y)
+        return {'loss': loss, 'log': {'train_loss': loss}}
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=0.02)
+
+
+@pytest.mark.parametrize("cls", [
+    SimpleNoArgsModel,
+    NoArgsSubClassEvalModel,
+])
+def test_model_nohparams_train_test(tmpdir, cls):
+    """Test models that do not tae any argument in init."""
+
+    model = cls()
+    trainer = Trainer(
+        max_epochs=1,
+        default_root_dir=tmpdir,
+    )
+
+    train_loader = DataLoader(TrialMNIST(os.getcwd(), train=True, download=True), batch_size=32)
+    trainer.fit(model, train_loader)
+
+    test_loader = DataLoader(TrialMNIST(os.getcwd(), train=False, download=True), batch_size=32)
+    trainer.test(test_dataloaders=test_loader)
+
+
+def test_model_ignores_non_exist_kwargument(tmpdir):
+    """Test that the model takes only valid class arguments."""
+
+    class LocalModel(EvalModelTemplate):
+        def __init__(self, batch_size=15):
+            super().__init__(batch_size=batch_size)
+            self.save_hyperparameters()
+
+    model = LocalModel()
+    assert model.hparams.batch_size == 15
+
+    # verify that the checkpoint saved the correct values
+    trainer = Trainer(default_root_dir=tmpdir, max_epochs=1)
+    trainer.fit(model)
+
+    # verify that we can overwrite whatever we want
+    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