hotfix to unblock hparams and OmniConf - removes auto_register_init_args by default (#2025)

* ogc install

* cleaned up tests

* hot fix

* hot fix

* hot fix

* hot fix

* hot fix

* hot fix

* hot fix

* hot fix

* hot fix

pytorch_lightning/core/lightning.py

@@ -33,7 +33,7 @@ else:
 CHECKPOINT_KEY_MODULE_ARGS = 'module_arguments'
 
 
-class LightningModule(ABC, DeviceDtypeModuleMixin, GradInformation, ModelIO, ModelHooks):
+class LightningModule(ABC, DeviceDtypeModuleMixin, GradInformation, ModelIO, ModelHooks, torch.nn.Module):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -76,9 +76,6 @@ class LightningModule(ABC, DeviceDtypeModuleMixin, GradInformation, ModelIO, Mod
         #: device reference
         self._device = torch.device('cpu')
 
-        # register all params passed into the child module in __init__
-        self._auto_collect_arguments()
-
     @property
     def on_gpu(self):
         """
@@ -1701,7 +1698,7 @@ class LightningModule(ABC, DeviceDtypeModuleMixin, GradInformation, ModelIO, Mod
                        " and this method will be removed in v1.0.0", DeprecationWarning)
         return self.get_progress_bar_dict()
 
-    def _auto_collect_arguments(self):
+    def auto_collect_arguments(self):
         """Collect all arguments module arguments."""
         frame = inspect.currentframe()
 
@@ -1717,9 +1714,13 @@ class LightningModule(ABC, DeviceDtypeModuleMixin, GradInformation, ModelIO, Mod
     @property
     def module_arguments(self) -> dict:
         """Aggregate this module and all parents arguments."""
-        args = dict(self._module_parents_arguments)
-        args.update(self._module_self_arguments)
-        return args
+        try:
+            args = dict(self._module_parents_arguments)
+            args.update(self._module_self_arguments)
+            return args
+        except AttributeError as e:
+            rank_zero_warn('you called `module.module_arguments` without calling self.auto_collect_arguments()')
+            return {}
 
 
 def _collect_init_args(frame, path_args: list) -> list:

requirements-extra.txt

@@ -8,4 +8,5 @@ wandb>=0.8.21
 trains>=0.14.1
 matplotlib>=3.1.1
 # no need to install with [pytorch] as pytorch is already installed and torchvision is required only for Horovod examples
-horovod>=0.19.1
+horovod>=0.19.1
+omegaconf==2.0.0

tests/base/model_template.py

@@ -53,6 +53,8 @@ class EvalModelTemplate(
                  **kwargs) -> object:
         # init superclass
         super().__init__()
+        self.auto_collect_arguments()
+
         self.drop_prob = drop_prob
         self.batch_size = batch_size
         self.in_features = in_features

tests/base/model_train_steps.py

@@ -23,16 +23,12 @@ class TrainingStepVariations(ABC):
         loss_val = self.loss(y, y_hat)
 
         # alternate possible outputs to test
-        if self.trainer.batch_idx % 1 == 0:
-            output = OrderedDict({
-                'loss': loss_val,
-                'progress_bar': {'some_val': loss_val * loss_val},
-                'log': {'train_some_val': loss_val * loss_val},
-            })
-            return output
-
-        if self.trainer.batch_idx % 2 == 0:
-            return loss_val
+        output = OrderedDict({
+            'loss': loss_val,
+            'progress_bar': {'some_val': loss_val * loss_val},
+            'log': {'train_some_val': loss_val * loss_val},
+        })
+        return output
 
     def training_step__inf_loss(self, batch, batch_idx, optimizer_idx=None):
         output = self.training_step(batch, batch_idx, optimizer_idx)

tests/base/model_valid_epoch_ends.py

@@ -20,28 +20,35 @@ class ValidationEpochEndVariations(ABC):
             # recursive mean for multilevel dicts
             return torch.stack([x[key] if isinstance(x, dict) else _mean(x, key) for x in res]).mean()
 
-        # return torch.stack(outputs).mean()
         val_loss_mean = _mean(outputs, 'val_loss')
         val_acc_mean = _mean(outputs, 'val_acc')
-        for output in outputs:
-            val_loss = self.get_output_metric(output, 'val_loss')
-
-            # reduce manually when using dp
-            if self.trainer.use_dp or self.trainer.use_ddp2:
-                val_loss = torch.mean(val_loss)
-            val_loss_mean += val_loss
-
-            # reduce manually when using dp
-            val_acc = self.get_output_metric(output, 'val_acc')
-            if self.trainer.use_dp or self.trainer.use_ddp2:
-                val_acc = torch.mean(val_acc)
-
-            val_acc_mean += val_acc
-
-        if outputs:  # skip zero divisions
-            val_loss_mean /= len(outputs)
-            val_acc_mean /= len(outputs)
 
         metrics_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
         results = {'progress_bar': metrics_dict, 'log': metrics_dict}
         return results
+
+    def validation_epoch_end_multiple_dataloaders(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+
+        Args:
+            outputs: list of individual outputs of each validation step
+        """
+
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        def _mean(res, key):
+            return torch.stack([x[key] for x in res]).mean()
+
+        pbar = {}
+        logs = {}
+        for dl_output_list in outputs:
+            output_keys = dl_output_list[0].keys()
+            output_keys = [x for x in output_keys if 'val_' in x]
+            for key in output_keys:
+                metric_out = _mean(dl_output_list, key)
+                pbar[key] = metric_out
+                logs[key] = metric_out
+
+        results = {'progress_bar': pbar, 'log': logs}
+        return results

tests/base/model_valid_steps.py

@@ -23,33 +23,14 @@ class ValidationStepVariations(ABC):
         # acc
         labels_hat = torch.argmax(y_hat, dim=1)
         val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-        val_acc = torch.tensor(val_acc)
+        val_acc = torch.tensor(val_acc).type_as(x)
 
-        if self.on_gpu:
-            val_acc = val_acc.cuda(loss_val.device.index)
-
-        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
-        if self.trainer.use_dp:
-            loss_val = loss_val.unsqueeze(0)
-            val_acc = val_acc.unsqueeze(0)
-
-        # alternate possible outputs to test
-        if batch_idx % 1 == 0:
-            output = OrderedDict({
-                'val_loss': loss_val,
-                'val_acc': val_acc,
-            })
-            return output
-        if batch_idx % 2 == 0:
-            return val_acc
-
-        if batch_idx % 3 == 0:
-            output = OrderedDict({
-                'val_loss': loss_val,
-                'val_acc': val_acc,
-                'test_dic': {'val_loss_a': loss_val}
-            })
-            return output
+        output = OrderedDict({
+            'val_loss': loss_val,
+            'val_acc': val_acc,
+            'test_dic': {'val_loss_a': loss_val}
+        })
+        return output
 
     def validation_step__multiple_dataloaders(self, batch, batch_idx, dataloader_idx, **kwargs):
         """
@@ -66,36 +47,10 @@ class ValidationStepVariations(ABC):
         # acc
         labels_hat = torch.argmax(y_hat, dim=1)
         val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-        val_acc = torch.tensor(val_acc)
+        val_acc = torch.tensor(val_acc).type_as(x)
 
-        if self.on_gpu:
-            val_acc = val_acc.cuda(loss_val.device.index)
-
-        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
-        if self.trainer.use_dp:
-            loss_val = loss_val.unsqueeze(0)
-            val_acc = val_acc.unsqueeze(0)
-
-        # alternate possible outputs to test
-        if batch_idx % 1 == 0:
-            output = OrderedDict({
-                'val_loss': loss_val,
-                'val_acc': val_acc,
-            })
-            return output
-        if batch_idx % 2 == 0:
-            return val_acc
-
-        if batch_idx % 3 == 0:
-            output = OrderedDict({
-                'val_loss': loss_val,
-                'val_acc': val_acc,
-                'test_dic': {'val_loss_a': loss_val}
-            })
-            return output
-        if batch_idx % 5 == 0:
-            output = OrderedDict({
-                f'val_loss_{dataloader_idx}': loss_val,
-                f'val_acc_{dataloader_idx}': val_acc,
-            })
-            return output
+        output = OrderedDict({
+            f'val_loss_{dataloader_idx}': loss_val,
+            f'val_acc_{dataloader_idx}': val_acc,
+        })
+        return output

tests/models/test_hparams.py

@@ -3,9 +3,65 @@ import os
 import pytest
 import torch
 
-from pytorch_lightning import Trainer
+from pytorch_lightning import Trainer, LightningModule
 from pytorch_lightning.core.lightning import CHECKPOINT_KEY_MODULE_ARGS
 from tests.base import EvalModelTemplate
+from omegaconf import OmegaConf
+import sys
+
+
+class OmegaConfModel(EvalModelTemplate):
+    def __init__(self, ogc):
+        super().__init__()
+        self.ogc = ogc
+        self.size = ogc.list[0]
+
+
+def test_class_nesting(tmpdir):
+
+    class Module(LightningModule):
+        def forward(self):
+            return 0
+
+    # make sure PL modules are always nn.Module
+    a = Module()
+    assert isinstance(a, torch.nn.Module)
+
+    def test_outside():
+        a = Module()
+        print(a.module_arguments)
+
+    class A:
+        def test(self):
+            a = Module()
+            print(a.module_arguments)
+
+        def test2(self):
+            test_outside()
+
+    test_outside()
+    A().test2()
+    A().test()
+
+
+def test_omegaconf(tmpdir):
+
+    # ogc only for 3.8
+    major = sys.version_info[0]
+    minor = sys.version_info[1]
+    if major < 3 and minor < 8:
+        return
+
+    conf = OmegaConf.create({"k": "v", "list": [15.4, {"a": "1", "b": "2"}]})
+    model = OmegaConfModel(conf)
+
+    # ensure ogc passed values correctly
+    assert model.size == 15.4
+
+    trainer = Trainer(default_root_dir=tmpdir, max_epochs=2, overfit_pct=0.5)
+    result = trainer.fit(model)
+
+    assert result == 1
 
 
 class SubClassEvalModel(EvalModelTemplate):
@@ -14,6 +70,7 @@ class SubClassEvalModel(EvalModelTemplate):
     def __init__(self, *args, subclass_arg=1200, **kwargs):
         super().__init__(*args, **kwargs)
         self.subclass_arg = subclass_arg
+        self.auto_collect_arguments()
 
 
 class SubSubClassEvalModel(SubClassEvalModel):
@@ -25,6 +82,7 @@ class AggSubClassEvalModel(SubClassEvalModel):
     def __init__(self, *args, my_loss=torch.nn.CrossEntropyLoss(), **kwargs):
         super().__init__(*args, **kwargs)
         self.my_loss = my_loss
+        self.auto_collect_arguments()
 
 
 @pytest.mark.parametrize("cls", [EvalModelTemplate,
@@ -47,7 +105,7 @@ def test_collect_init_arguments(tmpdir, cls):
         assert isinstance(model.my_loss, torch.nn.CosineEmbeddingLoss)
 
     # verify that the checkpoint saved the correct values
-    trainer = Trainer(max_steps=5, default_root_dir=tmpdir)
+    trainer = Trainer(default_root_dir=tmpdir, max_epochs=2, overfit_pct=0.5)
     trainer.fit(model)
     raw_checkpoint_path = os.listdir(trainer.checkpoint_callback.dirpath)
     raw_checkpoint_path = [x for x in raw_checkpoint_path if '.ckpt' in x][0]

tests/trainer/test_dataloaders.py

@@ -74,6 +74,7 @@ def test_multiple_val_dataloader(tmpdir):
     model = EvalModelTemplate()
     model.val_dataloader = model.val_dataloader__multiple
     model.validation_step = model.validation_step__multiple_dataloaders
+    model.validation_epoch_end = model.validation_epoch_end_multiple_dataloaders
 
     # fit model
     trainer = Trainer(
@@ -193,6 +194,7 @@ def test_multiple_dataloaders_passed_to_fit(tmpdir):
 
     model = EvalModelTemplate()
     model.validation_step = model.validation_step__multiple_dataloaders
+    model.validation_epoch_end = model.validation_epoch_end_multiple_dataloaders
     model.test_step = model.test_step__multiple_dataloaders
 
     # train, multiple val and multiple test passed to fit