hparams as dict [blocked by 1041] (#1029)

* hparams as dict

* hparams as dict

* fixing

* fixing

* fixing

* fixing

* typing

* typing

* chnagelog

* update set hparams

* use setter

* simplify

* chnagelog

* imports

* pylint

* typing

* Update training_io.py

* Update training_io.py

* Update lightning.py

* Update test_trainer.py

* Update __init__.py

* Update base.py

* Update utils.py

* Update test_trainer.py

* Update training_io.py

* Update test_trainer.py

* Update test_trainer.py

* Update test_trainer.py

* Update test_trainer.py

* Update callback_config.py

* Update callback_config.py

* Update test_trainer.py

Co-authored-by: William Falcon <waf2107@columbia.edu>

tests/loggers/test_wandb.py

@@ -24,8 +24,8 @@ def test_wandb_logger(wandb):
     logger.log_metrics({'acc': 1.0}, step=3)
     wandb.init().log.assert_called_once_with({'global_step': 3, 'acc': 1.0})
 
-    logger.log_hyperparams('test')
-    wandb.init().config.update.assert_called_once_with('test')
+    logger.log_hyperparams({'test': None})
+    wandb.init().config.update.assert_called_once_with({'test': None})
 
     logger.watch('model', 'log', 10)
     wandb.watch.assert_called_once_with('model', log='log', log_freq=10)

tests/models/__init__.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from .base import TestModelBase
+from .base import TestModelBase, DictHparamsModel
 from .mixins import (
     LightEmptyTestStep,
     LightValidationStepMixin,

tests/models/base.py

@@ -6,9 +6,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import optim
 from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
 from torchvision import transforms
 from torchvision.datasets import MNIST
+from typing import Dict
 
 try:
     from test_tube import HyperOptArgumentParser
@@ -36,6 +36,28 @@ class TestingMNIST(MNIST):
         self.targets = self.targets[:num_samples]
 
 
+class DictHparamsModel(LightningModule):
+
+    def __init__(self, hparams: Dict):
+        super(DictHparamsModel, self).__init__()
+        self.l1 = torch.nn.Linear(hparams.get('in_features'), hparams['out_features'])
+
+    def forward(self, x):
+        return torch.relu(self.l1(x.view(x.size(0), -1)))
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'loss': F.cross_entropy(y_hat, y)}
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=0.02)
+
+    def train_dataloader(self):
+        return DataLoader(MNIST(os.getcwd(), train=True, download=True,
+                                transform=transforms.ToTensor()), batch_size=32)
+
+
 class TestModelBase(LightningModule):
     """
     Base LightningModule for testing. Implements only the required

tests/models/utils.py

@@ -168,6 +168,20 @@ def load_model(exp, root_weights_dir, module_class=LightningTemplateModel, path_
     return trained_model
 
 
+def load_model_from_checkpoint(root_weights_dir, module_class=LightningTemplateModel):
+    # load trained model
+    checkpoints = [x for x in os.listdir(root_weights_dir) if '.ckpt' in x]
+    weights_dir = os.path.join(root_weights_dir, checkpoints[0])
+
+    trained_model = module_class.load_from_checkpoint(
+        checkpoint_path=weights_dir,
+    )
+
+    assert trained_model is not None, 'loading model failed'
+
+    return trained_model
+
+
 def run_prediction(dataloader, trained_model, dp=False, min_acc=0.45):
     # run prediction on 1 batch
     for batch in dataloader:

tests/trainer/test_trainer.py

@@ -3,30 +3,28 @@ import math
 import os
 import pytest
 import torch
-import argparse
+from argparse import ArgumentParser, Namespace
 
 import tests.models.utils as tutils
 from unittest import mock
-from pytorch_lightning import Trainer
+from pytorch_lightning import Trainer, LightningModule
 from pytorch_lightning.callbacks import (
     EarlyStopping,
     ModelCheckpoint,
 )
 from tests.models import (
     TestModelBase,
+    DictHparamsModel,
     LightningTestModel,
     LightEmptyTestStep,
     LightValidationStepMixin,
     LightValidationMultipleDataloadersMixin,
     LightTrainDataloader,
     LightTestDataloader,
-    LightValidationMixin,
-    LightTestMixin
 )
 from pytorch_lightning.core.lightning import load_hparams_from_tags_csv
 from pytorch_lightning.trainer.logging import TrainerLoggingMixin
 from pytorch_lightning.utilities.debugging import MisconfigurationException
-from pytorch_lightning import Callback
 
 
 def test_no_val_module(tmpdir):
@@ -128,7 +126,7 @@ def test_gradient_accumulation_scheduling(tmpdir):
         assert Trainer(accumulate_grad_batches={1: 2.5, 3: 5})
 
     # test optimizer call freq matches scheduler
-    def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, second_order_closure=None):
+    def _optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, second_order_closure=None):
         # only test the first 12 batches in epoch
         if batch_idx < 12:
             if epoch == 0:
@@ -179,7 +177,7 @@ def test_gradient_accumulation_scheduling(tmpdir):
                       default_save_path=tmpdir)
 
     # for the test
-    trainer.optimizer_step = optimizer_step
+    trainer.optimizer_step = _optimizer_step
     model.prev_called_batch_idx = 0
 
     trainer.fit(model)
@@ -188,7 +186,6 @@ def test_gradient_accumulation_scheduling(tmpdir):
 def test_loading_meta_tags(tmpdir):
     tutils.reset_seed()
 
-    from argparse import Namespace
     hparams = tutils.get_hparams()
 
     # save tags
@@ -604,8 +601,9 @@ def test_testpass_overrides(tmpdir):
     model = LightningTestModel(hparams)
     Trainer().test(model)
 
+
 @mock.patch('argparse.ArgumentParser.parse_args',
-            return_value=argparse.Namespace(**Trainer.default_attributes()))
+            return_value=Namespace(**Trainer.default_attributes()))
 def test_default_args(tmpdir):
     """Tests default argument parser for Trainer"""
     tutils.reset_seed()
@@ -613,7 +611,7 @@ def test_default_args(tmpdir):
     # logger file to get meta
     logger = tutils.get_test_tube_logger(tmpdir, False)
 
-    parser = argparse.ArgumentParser(add_help=False)
+    parser = ArgumentParser(add_help=False)
     args = parser.parse_args()
     args.logger = logger
 
@@ -622,3 +620,25 @@ def test_default_args(tmpdir):
 
     assert isinstance(trainer, Trainer)
     assert trainer.max_epochs == 5
+
+
+def test_hparams_save_load(tmpdir):
+    model = DictHparamsModel({'in_features': 28 * 28, 'out_features': 10})
+
+    # logger file to get meta
+    trainer_options = dict(
+        default_save_path=tmpdir,
+        max_epochs=2,
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    assert result == 1
+
+    # try to load the model now
+    pretrained_model = tutils.load_model_from_checkpoint(
+        trainer.checkpoint_callback.dirpath,
+        module_class=DictHparamsModel
+    )