diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py index fab1d3a8..fb3e2d0e 100644 --- a/pytorch_lightning/core/lightning.py +++ b/pytorch_lightning/core/lightning.py @@ -1,4 +1,5 @@ import collections +import inspect import logging as log import csv import os @@ -15,6 +16,7 @@ from pytorch_lightning.core.hooks import ModelHooks from pytorch_lightning.core.saving import ModelIO from pytorch_lightning.core.memory import ModelSummary from pytorch_lightning.overrides.data_parallel import LightningDistributedDataParallel +from pytorch_lightning.utilities.debugging import MisconfigurationException try: import torch_xla.core.xla_model as xm @@ -1111,13 +1113,10 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): else: checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage) - # 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) + # add the hparams from csv file to checkpoint + checkpoint['hparams'] = vars(hparams) + model = cls._load_model_state(checkpoint) return model @classmethod @@ -1182,17 +1181,36 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): else: 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) + model = cls._load_model_state(checkpoint) + return model + + @classmethod + def _load_model_state(cls, checkpoint): + cls_takes_hparams = 'hparams' in inspect.signature(cls.__init__).parameters + ckpt_hparams = checkpoint.get('hparams') + + if cls_takes_hparams: + if ckpt_hparams is not None: + hparams = Namespace(**ckpt_hparams) + else: + warnings.warn( + f"Checkpoint does not contain hyperparameters but {cls.__name__}'s __init__ contains" + " argument 'hparams'. Will pass in an empty Namespace instead." + " Did you forget to store your model hyperparameters in self.hparams?" + ) + hparams = Namespace() + else: # The user's LightningModule does not define a hparams argument + if ckpt_hparams is None: + hparams = None + else: + raise MisconfigurationException( + f"Checkpoint contains hyperparameters but {cls.__name__}'s __init__ is missing the" + " argument 'hparams'. Are you loading the correct checkpoint?" + ) # load the state_dict on the model automatically - model = cls(hparams) + model_args = [hparams] if hparams else [] + model = cls(*model_args) model.load_state_dict(checkpoint['state_dict']) # give model a chance to load something diff --git a/tests/models/__init__.py b/tests/models/__init__.py index df16bffd..ad926b60 100644 --- a/tests/models/__init__.py +++ b/tests/models/__init__.py @@ -26,3 +26,21 @@ class LightningTestModel(LightningValidationMixin, LightningTestMixin, Lightning def on_training_metrics(self, logs): logs['some_tensor_to_test'] = torch.rand(1) + + +class LightningTestModelWithoutHyperparametersArg(LightningTestModel): + """ without hparams argument in constructor """ + + def __init__(self): + import tests.models.utils as tutils + + # the user loads the hparams in some other way + hparams = tutils.get_hparams() + super().__init__(hparams) + + +class LightningTestModelWithUnusedHyperparametersArg(LightningTestModelWithoutHyperparametersArg): + """ has hparams argument in constructor but is not used """ + + def __init__(self, hparams): + super().__init__() diff --git a/tests/test_restore_models.py b/tests/test_restore_models.py index ba347e47..1ed36c2a 100644 --- a/tests/test_restore_models.py +++ b/tests/test_restore_models.py @@ -1,12 +1,18 @@ import logging as log import os +import pytest import torch import tests.models.utils as tutils from pytorch_lightning import Trainer from pytorch_lightning.callbacks import ModelCheckpoint -from tests.models import LightningTestModel +from pytorch_lightning.utilities.debugging import MisconfigurationException +from tests.models import ( + LightningTestModel, + LightningTestModelWithoutHyperparametersArg, + LightningTestModelWithUnusedHyperparametersArg +) def test_running_test_pretrained_model_ddp(tmpdir): @@ -380,5 +386,39 @@ def test_model_saving_loading(tmpdir): new_pred = model_2(x) assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1 + +def test_load_model_with_missing_hparams(tmpdir): + trainer_options = dict( + show_progress_bar=False, + max_epochs=1, + checkpoint_callback=ModelCheckpoint(tmpdir, save_top_k=-1), + logger=False, + default_save_path=tmpdir, + ) + + # fit model + trainer = Trainer(**trainer_options) + + model = LightningTestModelWithoutHyperparametersArg() + trainer.fit(model) + last_checkpoint = os.path.join(trainer.checkpoint_callback.filepath, "_ckpt_epoch_0.ckpt") + + # try to load a checkpoint that has hparams but model is missing hparams arg + with pytest.raises(MisconfigurationException, match=r".*__init__ is missing the argument 'hparams'.*"): + LightningTestModelWithoutHyperparametersArg.load_from_checkpoint(last_checkpoint) + + # create a checkpoint without hyperparameters + # if the model does not take a hparams argument, it should not throw an error + ckpt = torch.load(last_checkpoint) + del(ckpt['hparams']) + torch.save(ckpt, last_checkpoint) + LightningTestModelWithoutHyperparametersArg.load_from_checkpoint(last_checkpoint) + + # load checkpoint without hparams again + # warn if user's model has hparams argument + with pytest.warns(UserWarning, match=r".*Will pass in an empty Namespace instead."): + LightningTestModelWithUnusedHyperparametersArg.load_from_checkpoint(last_checkpoint) + + # if __name__ == '__main__': # pytest.main([__file__])