diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py index 0e9d00c6..4bb3c406 100644 --- a/pytorch_lightning/trainer/training_io.py +++ b/pytorch_lightning/trainer/training_io.py @@ -325,6 +325,7 @@ class TrainerIOMixin(ABC): checkpoint['native_amp_scaling_state'] = self.scaler.state_dict() if hasattr(model, "hparams"): + self.__clean_namespace(model.hparams) is_namespace = isinstance(model.hparams, Namespace) checkpoint['hparams'] = vars(model.hparams) if is_namespace else model.hparams checkpoint['hparams_type'] = 'namespace' if is_namespace else 'dict' @@ -338,6 +339,31 @@ class TrainerIOMixin(ABC): return checkpoint + def __clean_namespace(self, hparams): + """ + Removes all functions from hparams so we can pickle + :param hparams: + :return: + """ + + if isinstance(hparams, Namespace): + del_attrs = [] + for k in hparams.__dict__: + if callable(getattr(hparams, k)): + del_attrs.append(k) + + for k in del_attrs: + delattr(hparams, k) + + elif isinstance(hparams, dict): + del_attrs = [] + for k, v in hparams.items(): + if callable(v): + del_attrs.append(k) + + for k in del_attrs: + del hparams[k] + # -------------------- # HPC IO # -------------------- diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py index 2d34be24..6876a693 100644 --- a/tests/trainer/test_trainer.py +++ b/tests/trainer/test_trainer.py @@ -27,7 +27,7 @@ from tests.base import ( def test_hparams_save_load(tmpdir): - model = DictHparamsModel({'in_features': 28 * 28, 'out_features': 10}) + model = DictHparamsModel({'in_features': 28 * 28, 'out_features': 10, 'failed_key': lambda x: x}) # logger file to get meta trainer_options = dict( @@ -79,12 +79,13 @@ def test_no_val_module(tmpdir): new_weights_path = os.path.join(tmpdir, 'save_test.ckpt') trainer.save_checkpoint(new_weights_path) - # load new model - tags_path = tutils.get_data_path(logger, path_dir=tmpdir) - tags_path = os.path.join(tags_path, 'meta_tags.csv') + # assert ckpt has hparams + ckpt = torch.load(new_weights_path) + assert 'hparams' in ckpt.keys(), 'hparams missing from checkpoints' + + # won't load without hparams in the ckpt model_2 = LightningTestModel.load_from_checkpoint( checkpoint_path=new_weights_path, - tags_csv=tags_path ) model_2.eval()