fix hparams issue (#1623)

fix hparams issue

fix hparams issue
This commit is contained in:
William Falcon
2020-04-27 15:48:58 +02:00
committed by J. Borovec
parent 813e37916d
commit afd0e5d02f
2 changed files with 32 additions and 5 deletions
+26
View File
@@ -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
# --------------------
+6 -5
View File
@@ -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()