Merge pull request #1623 from PyTorchLightning/hparams

Hparams

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
     # --------------------

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()