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/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: