Repository: Lightning-AI/pytorch-lightning

Issue:
load_from_checkpoint() doesn't work when a LightningModule inherits from typing.Generic
## 🐛 Bug

When a LightningModule with saved hyperparameters inherits from `typing.Generic`, hyperparameters saved in the checkpoint file are not loaded automatically, causing an error. When `load_from_checkpoint()` calls `inspect.signature()` to gather the list of arguments of the LightningModule that inherits from `typing.Generic`, `inspect.signature()` returns `['args', 'kwds']` instead of the actual arguments, because `typing.Generic` implements an empty `__new__()` (the execution path ends up here: https://github.com/python/cpython/blob/3.8/Lib/inspect.py#L2324). As a result, PL filters out all the saved hyperparameters from the checkpoint, which results in an error when trying to instantiate the LightningModule. I'd assume this would happen when a LightningModule inherits from any class that implements `__new__()` such as `abc.ABC`.

### To Reproduce
Create a LightningModule that inherits from `typing.Generic` with some hyperparameters, fit it, then try to load it from a checkpoint.


#### Code sample
```python
import torch
import torch.nn.functional as F
import pytorch_lightning as pl

from typing import Generic, TypeVar
from torch.utils.data import DataLoader

T = TypeVar("T")


class GenericLitClassifier(Generic[T], pl.LightningModule):
    def __init__(self, dim):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, dim)
        self.save_hyperparameters()

    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 configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


class LitClassifier(GenericLitClassifier[str]):
    pass


class Dataset:
    def __getitem__(self, idx):
        return torch.ones(1, 784), 1

    def __len__(self):
        return 5


train_loader = DataLoader(Dataset(), batch_size=2)
model = LitClassifier(10)
trainer = pl.Trainer(max_epochs=5)
trainer.fit(model, train_loader)

for path, _ in trainer.checkpoint_callback.best_k_models.items():
    lm = LitClassifier.load_from_checkpoint(path)
```

### Expected behavior

Even when a LightningModule inherits from any class that implements `__new__()` (e.g. `typing.Generic`), its hyperparameters should be loaded automatically from a checkpoint.


### Environment

```
* CUDA:
        - GPU:
        - available:         False
        - version:           None
* Packages:
        - numpy:             1.19.1
        - pyTorch_debug:     False
        - pyTorch_version:   1.5.1
        - pytorch-lightning: 0.8.5
        - tensorboard:       2.3.0
        - tqdm:              4.48.2
* System:
        - OS:                Darwin
        - architecture:
                - 64bit
                - 
        - processor:         i386
        - python:            3.7.7
        - version:           Darwin Kernel Version 18.7.0: Mon Apr 27 20:09:39 PDT 2020; root:xnu-4903.278.35~1/RELEASE_X86_64
```

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:

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
@@ -48,13 +48,13 @@ class ModelIO(object):
 
     @classmethod
     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,
-            **kwargs
+        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,
+        **kwargs,
     ):
         r"""
         Primary way of loading a model from a checkpoint. When Lightning saves a checkpoint
@@ -156,7 +156,7 @@ def load_from_checkpoint(
     @classmethod
     def _load_model_state(cls, checkpoint: Dict[str, Any], strict: bool = True, *cls_args, **cls_kwargs):
         cls_spec = inspect.getfullargspec(cls.__init__)
-        cls_init_args_name = inspect.signature(cls).parameters.keys()
+        cls_init_args_name = inspect.signature(cls.__init__).parameters.keys()
         # pass in the values we saved automatically
         if cls.CHECKPOINT_HYPER_PARAMS_KEY in checkpoint:
             model_args = {}


Test patch:
diff --git a/tests/base/__init__.py b/tests/base/__init__.py
--- a/tests/base/__init__.py
+++ b/tests/base/__init__.py
@@ -1,4 +1,4 @@
 """Models for testing."""
 
 from tests.base.datasets import TrialMNIST
-from tests.base.model_template import EvalModelTemplate
+from tests.base.model_template import EvalModelTemplate, GenericEvalModelTemplate
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
@@ -2,6 +2,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from typing import Generic, TypeVar
+
 from pytorch_lightning.core.lightning import LightningModule
 from tests.base.datasets import TrialMNIST, PATH_DATASETS
 from tests.base.model_optimizers import ConfigureOptimizersPool
@@ -28,7 +30,7 @@ class EvalModelTemplate(
     ValDataloaderVariations,
     TestDataloaderVariations,
     ConfigureOptimizersPool,
-    LightningModule
+    LightningModule,
 ):
     """
     This template houses all  combinations of model  configurations  we want to test
@@ -37,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__()
@@ -83,17 +85,11 @@ def __build_model(self):
         Simple model for testing
         :return:
         """
-        self.c_d1 = nn.Linear(
-            in_features=self.in_features,
-            out_features=self.hidden_dim
-        )
+        self.c_d1 = nn.Linear(in_features=self.in_features, out_features=self.hidden_dim)
         self.c_d1_bn = nn.BatchNorm1d(self.hidden_dim)
         self.c_d1_drop = nn.Dropout(self.drop_prob)
 
-        self.c_d2 = nn.Linear(
-            in_features=self.hidden_dim,
-            out_features=self.out_features
-        )
+        self.c_d2 = nn.Linear(in_features=self.hidden_dim, out_features=self.out_features)
 
     def forward(self, x):
         x = self.c_d1(x)
@@ -130,8 +126,42 @@ def get_default_hparams(continue_training: bool = False, hpc_exp_number: int = 0
 
         if continue_training:
             args.update(
-                test_tube_do_checkpoint_load=True,
-                hpc_exp_number=hpc_exp_number,
+                test_tube_do_checkpoint_load=True, hpc_exp_number=hpc_exp_number,
             )
 
         return args
+
+
+T = TypeVar('T')
+
+
+class GenericParentEvalModelTemplate(Generic[T], EvalModelTemplate):
+    def __init__(
+        self,
+        drop_prob: float,
+        batch_size: int,
+        in_features: int,
+        learning_rate: float,
+        optimizer_name: str,
+        data_root: str,
+        out_features: int,
+        hidden_dim: int,
+        b1: float,
+        b2: float,
+    ):
+        super().__init__(
+            drop_prob=drop_prob,
+            batch_size=batch_size,
+            in_features=in_features,
+            learning_rate=learning_rate,
+            optimizer_name=optimizer_name,
+            data_root=data_root,
+            out_features=out_features,
+            hidden_dim=hidden_dim,
+            b1=b1,
+            b2=b2,
+        )
+
+
+class GenericEvalModelTemplate(GenericParentEvalModelTemplate[int]):
+    pass
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
@@ -12,7 +12,7 @@
 import tests.base.develop_utils as tutils
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
-from tests.base import EvalModelTemplate
+from tests.base import EvalModelTemplate, GenericEvalModelTemplate
 
 
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
@@ -151,10 +151,11 @@ def test_running_test_pretrained_model_cpu(tmpdir):
     tutils.assert_ok_model_acc(new_trainer)
 
 
-def test_load_model_from_checkpoint(tmpdir):
+@pytest.mark.parametrize('model_template', [EvalModelTemplate, GenericEvalModelTemplate])
+def test_load_model_from_checkpoint(tmpdir, model_template):
     """Verify test() on pretrained model."""
-    hparams = EvalModelTemplate.get_default_hparams()
-    model = EvalModelTemplate(**hparams)
+    hparams = model_template.get_default_hparams()
+    model = model_template(**hparams)
 
     trainer_options = dict(
         progress_bar_refresh_rate=0,
@@ -175,7 +176,7 @@ def test_load_model_from_checkpoint(tmpdir):
 
     # load last checkpoint
     last_checkpoint = sorted(glob.glob(os.path.join(trainer.checkpoint_callback.dirpath, "*.ckpt")))[-1]
-    pretrained_model = EvalModelTemplate.load_from_checkpoint(last_checkpoint)
+    pretrained_model = model_template.load_from_checkpoint(last_checkpoint)
 
     # test that hparams loaded correctly
     for k, v in hparams.items():
@@ -198,12 +199,7 @@ def test_dp_resume(tmpdir):
     hparams = EvalModelTemplate.get_default_hparams()
     model = EvalModelTemplate(**hparams)
 
-    trainer_options = dict(
-        max_epochs=1,
-        gpus=2,
-        distributed_backend='dp',
-        default_root_dir=tmpdir,
-    )
+    trainer_options = dict(max_epochs=1, gpus=2, distributed_backend='dp', default_root_dir=tmpdir,)
 
     # get logger
     logger = tutils.get_default_logger(tmpdir)
@@ -275,10 +271,7 @@ def test_model_saving_loading(tmpdir):
 
     # fit model
     trainer = Trainer(
-        max_epochs=1,
-        logger=logger,
-        checkpoint_callback=ModelCheckpoint(tmpdir),
-        default_root_dir=tmpdir,
+        max_epochs=1, logger=logger, checkpoint_callback=ModelCheckpoint(tmpdir), default_root_dir=tmpdir,
     )
     result = trainer.fit(model)
 
@@ -308,10 +301,7 @@ def test_model_saving_loading(tmpdir):
     # load new model
     hparams_path = tutils.get_data_path(logger, path_dir=tmpdir)
     hparams_path = os.path.join(hparams_path, 'hparams.yaml')
-    model_2 = EvalModelTemplate.load_from_checkpoint(
-        checkpoint_path=new_weights_path,
-        hparams_file=hparams_path,
-    )
+    model_2 = EvalModelTemplate.load_from_checkpoint(checkpoint_path=new_weights_path, hparams_file=hparams_path,)
     model_2.eval()
 
     # make prediction
@@ -335,10 +325,7 @@ def test_strict_model_load_more_params(monkeypatch, tmpdir, tmpdir_server, url_c
 
     # fit model
     trainer = Trainer(
-        default_root_dir=tmpdir,
-        max_epochs=1,
-        logger=logger,
-        checkpoint_callback=ModelCheckpoint(tmpdir),
+        default_root_dir=tmpdir, max_epochs=1, logger=logger, checkpoint_callback=ModelCheckpoint(tmpdir),
     )
     result = trainer.fit(model)
 
@@ -355,16 +342,12 @@ def test_strict_model_load_more_params(monkeypatch, tmpdir, tmpdir_server, url_c
     ckpt_path = hparams_url if url_ckpt else new_weights_path
 
     EvalModelTemplate.load_from_checkpoint(
-        checkpoint_path=ckpt_path,
-        hparams_file=hparams_path,
-        strict=False,
+        checkpoint_path=ckpt_path, hparams_file=hparams_path, strict=False,
     )
 
     with pytest.raises(RuntimeError, match=r'Unexpected key\(s\) in state_dict: "c_d3.weight", "c_d3.bias"'):
         EvalModelTemplate.load_from_checkpoint(
-            checkpoint_path=ckpt_path,
-            hparams_file=hparams_path,
-            strict=True,
+            checkpoint_path=ckpt_path, hparams_file=hparams_path, strict=True,
         )
 
 
@@ -381,10 +364,7 @@ def test_strict_model_load_less_params(monkeypatch, tmpdir, tmpdir_server, url_c
 
     # fit model
     trainer = Trainer(
-        default_root_dir=tmpdir,
-        max_epochs=1,
-        logger=logger,
-        checkpoint_callback=ModelCheckpoint(tmpdir),
+        default_root_dir=tmpdir, max_epochs=1, logger=logger, checkpoint_callback=ModelCheckpoint(tmpdir),
     )
     result = trainer.fit(model)
 
@@ -406,16 +386,12 @@ def __init__(self):
             self.c_d3 = torch.nn.Linear(7, 7)
 
     CurrentModel.load_from_checkpoint(
-        checkpoint_path=ckpt_path,
-        hparams_file=hparams_path,
-        strict=False,
+        checkpoint_path=ckpt_path, hparams_file=hparams_path, strict=False,
     )
 
     with pytest.raises(RuntimeError, match=r'Missing key\(s\) in state_dict: "c_d3.weight", "c_d3.bias"'):
         CurrentModel.load_from_checkpoint(
-            checkpoint_path=ckpt_path,
-            hparams_file=hparams_path,
-            strict=True,
+            checkpoint_path=ckpt_path, hparams_file=hparams_path, strict=True,
         )