Repository: Lightning-AI/pytorch-lightning
Pull request: update hparams, allow OmegaConf
## What does this PR do?
Fixes #2027
Fixes #2040
Fixes #2087

## PR review    
Anyone in the community is free to review the PR once the tests have passed.     
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Support DictConfig
We need to add DictConfig support for Omegaconf @Borda to the auto hparam save

Support omegaconf in hparams
We need to get hydra dict config to work.

Proposed solution:
    1. save all primitives to the state of the checkpoint.
    2. create function save_arguments()
        1. if no args passed- save everything in scope, even if it's not primitives.
        2. if kargs specified- save key value pairs.

Make sure to explicitly test the omegaconf (saving and loading).

[bug] remove frame introspection
Currently breaking master, creating pickling errors.
Instead saving params will have to be done explicitly.

Hints from issue discussion:
Yes, I agree. Users should have other options besides argument parser to set up their configuration. Personally speaking, I don't like typing too much on the command line b/c that is error-prone. A dictionary-like configuration system would be great. One example would be Ross's [yacs](https://github.com/rbgirshick/yacs) which works pretty well. 
OmegaConf is along the same lines as Yacs, but with more features (might as well support YACS too, but just saying).
@DKandrew @Darktex you would keep passing one eg `OmegaConf` argument which is used internally? Kind of pseudocode...
1. ```
    conf = OmegaConf(...)
    model = MyModel(conf)
    ```
2. ```
    conf = OmegaConf(...)
    model = MyModel(**vars(conf))
    ```
assuming that the conf can be also loaded from a file...
Hi @Borda 
I am not sure if I understand your question correctly, are you asking which approach you listed above is better? 
> I am not sure if I understand your question correctly, are you asking which approach you listed above is better?

I am trying to understand your use-case, mind draw your use case?
Sure. My use-case is the first one
```
conf = OmegaConf(...)
model = MyModel(conf)
```
I don't want to use the second case because `conf` contains too many entries: entries for network layers, dataloader, training, optimizer/scheduler, etc. I don't want to expand all of them into my `__init__()` because that would be too long. After all, MyModel is just an nn.Module with additional features, so internally I have a `self.conf` to store the entire config and use it whenever I need it. 

```
class MyModel(LightningModule)
    def __init__(self, conf):
        self.conf = conf
        # Normal network parameters like in_features, out_features
```

let’s do this:

## case 1
User explicitly says what they want to save.

```python
class LitModel(...):

    def __init__(self, conf):
        self.save_hyperparameters(conf)
```

## Case 2:
User wants to save all the init stuff.
They can do it all manually or ask us to do it automatically

```python
class LitModel(...):

    def __init__(self, arg1, arg2, arg3):
        # manually
        self.save_hyperparameters(arg_name=arg1, arg_name=arg2, arg_name=arg3)
        
        # equivalent automatic
        self.save_hyperparameters()
```

## Case 3:
They want to save ONLY some of the init stuff
```python
class LitModel(...):

    def __init__(self, arg1, arg2, arg3):
        # manually
        self.save_hyperparameters(arg_name=arg2)
```

## Special cases:
- namespace
```python
    def __init__(self, hparams):
        # manually
        self.save_hyperparameters(hparams)
```

- dict
```python
    def __init__(self, some_dict):
        # manually
        self.save_hyperparameters(some_dict)
```

- omniconf
```python
    def __init__(self, conf):
        # manually
        self.save_hyperparameters(conf)
```

- anything they want
```python
    def __init__(self, some_random_alternative_to_config):
        # manually
        self.save_hyperparameters(some_random_alternative_to_config)
```

@PyTorchLightning/core-contributors 

> let’s do this:

Sounds good to me. Looking forward to it! 
> ## case 1
> User explicitly says what they want to save.
> 
> ```python
> class LitModel(...):
> 
>     def __init__(self, conf):
>         self.save_hyperparameters(conf)
> ```

is very tricky as we would need to do some pairing from init frame and hope none of the `conf1, conf2, ...` has the same value so I would skip this case...

> ## Special cases:
> * namespace
> * dict
> 
> ```python
>     def __init__(self, conf):
>         # manually
>         self.save_hyperparameters(conf)
> ```

here you want to unroll all elements?
 assume you have `some_dict=dict(a=1, b=3)` then you would in fact do something similar like `self.save_hyperparameters(**some_dict)`

if fact 
```
def save_hyperparameters(**kwargs):
    ...
    for elms in kwargs:
        if isisntance(conf, (dict, OmegaConf)):
            self.save_hyperparameters(**conf)
        if isisntance(conf, Namespace):
            self.save_hyperparameters(**vars(conf))
       ...
```
Hi! thanks for your contribution!, great first issue!
are you sure about the name `omniconf` I found this but it is not probably it https://pypi.org/project/omniconf

EDIT: probably it should be this one, right?
`from omegaconf import OmegaConf, DictConfig`
probably the same as #2027
Yes, i guess it is. 
definitely not omniconf.

https://github.com/omry/omegaconf
Apologize for the confusion, fixed now! Thanks!

Commit message: wip

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3f8487cc742ac..6bac0ad2d4ba9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -78,7 +78,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed an issue where local variables were being collected into module_arguments ([#2048](https://github.com/PyTorchLightning/pytorch-lightning/pull/2048))
 
-- Fixed an issue with `auto_collect_arguments` collecting local variables that are not constructor arguments and not working for signatures that have the instance not named `self` ([#2048](https://github.com/PyTorchLightning/pytorch-lightning/pull/2048))
+- Fixed an issue with `_auto_collect_arguments` collecting local variables that are not constructor arguments and not working for signatures that have the instance not named `self` ([#2048](https://github.com/PyTorchLightning/pytorch-lightning/pull/2048))
 
 ## [0.7.6] - 2020-05-16
 
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index 282521334af85..4554f021dba0d 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -1,4 +1,5 @@
 import collections
+import copy
 import inspect
 import os
 import warnings
@@ -31,6 +32,15 @@
 else:
     XLA_AVAILABLE = True
 
+ALLOWED_CONFIG_TYPES = (dict, Namespace)
+try:
+    from omegaconf import DictConfig, OmegaConf
+except ImportError:
+    pass
+else:
+    ALLOWED_CONFIG_TYPES = ALLOWED_CONFIG_TYPES + (DictConfig, OmegaConf)
+
+
 CHECKPOINT_KEY_MODULE_ARGS = 'module_arguments'
 
 
@@ -77,6 +87,9 @@ def __init__(self, *args, **kwargs):
         #: device reference
         self._device = torch.device('cpu')
 
+        self._module_self_arguments = {}
+        self._module_parents_arguments = {}
+
     @property
     def on_gpu(self):
         """
@@ -1714,13 +1727,14 @@ def get_tqdm_dict(self) -> Dict[str, Union[int, str]]:
                        " and this method will be removed in v1.0.0", DeprecationWarning)
         return self.get_progress_bar_dict()
 
-    def auto_collect_arguments(self) -> None:
+    def _auto_collect_arguments(self, frame=None) -> None:
         """
         Collect all module arguments in the current constructor and all child constructors.
         The child constructors are all the ``__init__`` methods that reach the current class through
         (chained) ``super().__init__()`` calls.
         """
-        frame = inspect.currentframe()
+        if not frame:
+            frame = inspect.currentframe()
 
         frame_args = _collect_init_args(frame.f_back, [])
         self_arguments = frame_args[-1]
@@ -1739,19 +1753,70 @@ def module_arguments(self) -> dict:
         Aggregate of arguments passed to the constructor of this module and all parents.
 
         Return:
-            a dict in which the keys are the union of all argument names in the constructor and all
-            parent constructors, excluding `self`, `*args` and `**kwargs`.
+            custom object or dict in which the keys are the union of all argument names in the constructor
+            and all parent constructors, excluding `self`, `*args` and `**kwargs`.
         """
-        try:
-            args = dict(self._module_parents_arguments)
+        if isinstance(self._module_self_arguments, dict):
+            args = copy.deepcopy(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 {}
+        return copy.deepcopy(self._module_self_arguments)
+
+    def save_hyperparameters(self, *args, **kwargs) -> None:
+        """
+
+        >>> from collections import OrderedDict
+        >>> class ManuallyArgsModel(LightningModule):
+        ...     def __init__(self, arg1, arg2, arg3):
+        ...         super().__init__()
+        ...         # manually assin arguments
+        ...         self.save_hyperparameters(arg_name1=arg1, arg_name2=arg2, arg_name3=arg3)
+        ...     def forward(self, *args, **kwargs):
+        ...         ...
+        >>> model = ManuallyArgsModel(1, 'abc', 3.14)
+        >>> OrderedDict(model.module_arguments)
+        OrderedDict([('arg_name1', 1), ('arg_name2', 'abc'), ('arg_name3', 3.14)])
+
+        >>> from collections import OrderedDict
+        >>> class AutomaticArgsModel(LightningModule):
+        ...     def __init__(self, arg1, arg2, arg3):
+        ...         super().__init__()
+        ...         # equivalent automatic
+        ...         self.save_hyperparameters()
+        ...     def forward(self, *args, **kwargs):
+        ...         ...
+        >>> model = AutomaticArgsModel(1, 'abc', 3.14)
+        >>> OrderedDict(model.module_arguments)
+        OrderedDict([('arg1', 1), ('arg2', 'abc'), ('arg3', 3.14)])
+
+        >>> class SingleArgModel(LightningModule):
+        ...     def __init__(self, hparams):
+        ...         super().__init__()
+        ...         # manually assign single argument
+        ...         self.save_hyperparameters(hparams)
+        ...     def forward(self, *args, **kwargs):
+        ...         ...
+        >>> model = SingleArgModel(Namespace(p1=1, p2='abc', p3=3.14))
+        >>> model.module_arguments
+        Namespace(p1=1, p2='abc', p3=3.14)
+        """
+        if not args and not kwargs:
+            self._auto_collect_arguments()
+            return
+
+        if args:
+            if len(args) > 1:
+                raise ValueError('Only one argument can be passed.')
+            arg = args[0]
+            if not isinstance(arg, ALLOWED_CONFIG_TYPES):
+                raise ValueError(f'Unsupported argument type `{type(arg)}`.')
+            self._module_self_arguments = copy.copy(arg)
+
+        elif kwargs:
+            self._module_self_arguments = copy.deepcopy(kwargs)
 
 
-def _collect_init_args(frame, path_args: list) -> list:
+def _collect_init_args(frame, path_args: list, inside: bool = False) -> list:
     """
     Recursively collects the arguments passed to the child constructors in the inheritance tree.
 
@@ -1769,9 +1834,9 @@ def _collect_init_args(frame, path_args: list) -> list:
         cls = local_vars['__class__']
         spec = inspect.getfullargspec(cls.__init__)
         init_parameters = inspect.signature(cls.__init__).parameters
-        self_identifier = spec.args[0]        # "self" unless user renames it (always first arg)
-        varargs_identifier = spec.varargs     # by convention this is named "*args"
-        kwargs_identifier = spec.varkw        # by convention this is named "**kwargs"
+        self_identifier = spec.args[0]  # "self" unless user renames it (always first arg)
+        varargs_identifier = spec.varargs  # by convention this is named "*args"
+        kwargs_identifier = spec.varkw  # by convention this is named "**kwargs"
         exclude_argnames = (
             varargs_identifier, kwargs_identifier, self_identifier, '__class__', 'frame', 'frame_args'
         )
@@ -1783,6 +1848,8 @@ def _collect_init_args(frame, path_args: list) -> list:
 
         # recursive update
         path_args.append(local_args)
-        return _collect_init_args(frame.f_back, path_args)
+        return _collect_init_args(frame.f_back, path_args, inside=True)
+    elif not inside:
+        return _collect_init_args(frame.f_back, path_args, inside)
     else:
         return path_args
diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py
index fd0385cde4b34..c5bbb3ef4ec74 100644
--- a/pytorch_lightning/trainer/training_io.py
+++ b/pytorch_lightning/trainer/training_io.py
@@ -124,6 +124,12 @@
     list, tuple, set, dict,
     Namespace,  # for back compatibility
 )
+try:
+    from omegaconf import DictConfig, OmegaConf
+except ImportError:
+    pass
+else:
+    PRIMITIVE_TYPES = PRIMITIVE_TYPES + (DictConfig, OmegaConf)
 
 
 class TrainerIOMixin(ABC):
diff --git a/tests/base/model_template.py b/tests/base/model_template.py
index 1a5b1381fb171..98de113d80974 100644
--- a/tests/base/model_template.py
+++ b/tests/base/model_template.py
@@ -51,7 +51,7 @@ def __init__(self,
                  **kwargs) -> object:
         # init superclass
         super().__init__()
-        self.auto_collect_arguments()
+        self._auto_collect_arguments()
 
         self.drop_prob = drop_prob
         self.batch_size = batch_size
diff --git a/tests/models/test_hparams.py b/tests/models/test_hparams.py
index d1c28c9e8e39f..59cdf33c96ed8 100644
--- a/tests/models/test_hparams.py
+++ b/tests/models/test_hparams.py
@@ -1,24 +1,17 @@
 import os
 import sys
+from argparse import Namespace
 
 import pytest
 import torch
-from omegaconf import OmegaConf
-from packaging import version
+from omegaconf import OmegaConf, DictConfig
 
 from pytorch_lightning import Trainer, LightningModule
 from pytorch_lightning.core.lightning import CHECKPOINT_KEY_MODULE_ARGS
 from tests.base import EvalModelTemplate
 
 
-class OmegaConfModel(EvalModelTemplate):
-    def __init__(self, ogc):
-        super().__init__()
-        self.ogc = ogc
-        self.size = ogc.list[0]
-
-
-def test_class_nesting(tmpdir):
+def test_class_nesting():
 
     class MyModule(LightningModule):
         def forward(self):
@@ -47,6 +40,12 @@ def test2(self):
 
 @pytest.mark.xfail(sys.version_info >= (3, 8), reason='OmegaConf only for Python >= 3.8')
 def test_omegaconf(tmpdir):
+    class OmegaConfModel(EvalModelTemplate):
+        def __init__(self, ogc):
+            super().__init__()
+            self.ogc = ogc
+            self.size = ogc.list[0]
+
     conf = OmegaConf.create({"k": "v", "list": [15.4, {"a": "1", "b": "2"}]})
     model = OmegaConfModel(conf)
 
@@ -64,8 +63,11 @@ class SubClassEvalModel(EvalModelTemplate):
 
     def __init__(self, *args, subclass_arg=1200, **kwargs):
         super().__init__(*args, **kwargs)
-        self.subclass_arg = subclass_arg
-        self.auto_collect_arguments()
+        self.save_hyperparameters()
+
+
+class SubSubClassEvalModel(SubClassEvalModel):
+    pass
 
 
 class UnconventionalArgsEvalModel(EvalModelTemplate):
@@ -74,21 +76,21 @@ class UnconventionalArgsEvalModel(EvalModelTemplate):
     def __init__(obj, *more_args, other_arg=300, **more_kwargs):
         # intentionally named obj
         super().__init__(*more_args, **more_kwargs)
-        obj.other_arg = other_arg
+        obj.save_hyperparameters()
         other_arg = 321
-        obj.auto_collect_arguments()
-
-
-class SubSubClassEvalModel(SubClassEvalModel):
-    pass
 
 
 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()
+        self.save_hyperparameters()
+
+
+class DictConfSubClassEvalModel(SubClassEvalModel):
+    def __init__(self, *args, dict_conf=OmegaConf.create(dict(my_param='something')), **kwargs):
+        super().__init__(*args, **kwargs)
+        self.save_hyperparameters()
 
 
 @pytest.mark.parametrize("cls", [
@@ -97,10 +99,15 @@ def __init__(self, *args, my_loss=torch.nn.CrossEntropyLoss(), **kwargs):
     SubSubClassEvalModel,
     AggSubClassEvalModel,
     UnconventionalArgsEvalModel,
+    DictConfSubClassEvalModel,
 ])
 def test_collect_init_arguments(tmpdir, cls):
     """ Test that the model automatically saves the arguments passed into the constructor """
-    extra_args = dict(my_loss=torch.nn.CosineEmbeddingLoss()) if cls is AggSubClassEvalModel else {}
+    extra_args = {}
+    if cls is AggSubClassEvalModel:
+        extra_args.update(my_loss=torch.nn.CosineEmbeddingLoss())
+    elif cls is DictConfSubClassEvalModel:
+        extra_args.update(dict_conf=OmegaConf.create(dict(my_param='anything')))
 
     model = cls(**extra_args)
     assert model.batch_size == 32
@@ -116,9 +123,7 @@ def test_collect_init_arguments(tmpdir, cls):
     # verify that the checkpoint saved the correct values
     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]
-    raw_checkpoint_path = os.path.join(trainer.checkpoint_callback.dirpath, raw_checkpoint_path)
+    raw_checkpoint_path = _raw_checkpoint_path(trainer)
 
     raw_checkpoint = torch.load(raw_checkpoint_path)
     assert CHECKPOINT_KEY_MODULE_ARGS in raw_checkpoint
@@ -131,11 +136,22 @@ def test_collect_init_arguments(tmpdir, cls):
     if isinstance(model, AggSubClassEvalModel):
         assert isinstance(model.my_loss, torch.nn.CrossEntropyLoss)
 
+    if isinstance(model, DictConfSubClassEvalModel):
+        assert isinstance(model.dict_conf, DictConfig)
+        assert model.dict_conf == 'anything'
+
     # verify that we can overwrite whatever we want
     model = cls.load_from_checkpoint(raw_checkpoint_path, batch_size=99)
     assert model.batch_size == 99
 
 
+def _raw_checkpoint_path(trainer) -> str:
+    raw_checkpoint_paths = os.listdir(trainer.checkpoint_callback.dirpath)
+    raw_checkpoint_path = [x for x in raw_checkpoint_paths if '.ckpt' in x][0]
+    raw_checkpoint_path = os.path.join(trainer.checkpoint_callback.dirpath, raw_checkpoint_path)
+    return raw_checkpoint_path
+
+
 class LocalVariableModel1(EvalModelTemplate):
     """ This model has the super().__init__() call at the end. """
 
@@ -147,14 +163,14 @@ def __init__(self, arg1, arg2, *args, **kwargs):
 
 
 class LocalVariableModel2(EvalModelTemplate):
-    """ This model has the auto_collect_arguments() call at the end. """
+    """ This model has the _auto_collect_arguments() call at the end. """
 
     def __init__(self, arg1, arg2, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.argument1 = arg1  # arg2 intentionally not set
         arg1 = 'overwritten'
         local_var = 1234
-        self.auto_collect_arguments()  # this is intentionally here at the end
+        self._auto_collect_arguments()  # this is intentionally here at the end
 
 
 @pytest.mark.parametrize("cls", [
@@ -167,3 +183,62 @@ def test_collect_init_arguments_with_local_vars(cls):
     assert 'local_var' not in model.module_arguments
     assert model.module_arguments['arg1'] == 'overwritten'
     assert model.module_arguments['arg2'] == 2
+
+
+class NamespaceArgModel(EvalModelTemplate):
+    def __init__(self, hparams: Namespace):
+        super().__init__()
+        self.save_hyperparameters(hparams)
+
+
+class DictArgModel(EvalModelTemplate):
+    def __init__(self, some_dict: dict):
+        super().__init__()
+        self.save_hyperparameters(some_dict)
+
+
+class OmegaConfArgModel(EvalModelTemplate):
+    def __init__(self, conf: OmegaConf):
+        super().__init__()
+        self.save_hyperparameters(conf)
+
+
+@pytest.mark.parametrize("cls,config", [
+    (NamespaceArgModel, Namespace(my_arg=42)),
+    (DictArgModel, dict(my_arg=42)),
+    (OmegaConfArgModel, OmegaConf.create(dict(my_arg=42))),
+])
+def test_single_config_models(tmpdir, cls, config):
+    """ Test that the model automatically saves the arguments passed into the constructor """
+    model = cls(config)
+
+    # verify that the checkpoint saved the correct values
+    trainer = Trainer(default_root_dir=tmpdir, max_epochs=2, overfit_pct=0.5)
+    trainer.fit(model)
+
+    # verify that model loads correctly
+    raw_checkpoint_path = _raw_checkpoint_path(trainer)
+    model = cls.load_from_checkpoint(raw_checkpoint_path)
+    assert model.module_arguments == config
+
+
+class AnotherArgModel(EvalModelTemplate):
+    def __init__(self, arg1):
+        super().__init__()
+        self.save_hyperparameters(arg1)
+
+
+class OtherArgsModel(EvalModelTemplate):
+    def __init__(self, arg1, arg2):
+        super().__init__()
+        self.save_hyperparameters(arg1, arg2)
+
+
+@pytest.mark.parametrize("cls,config", [
+    (AnotherArgModel, dict(arg1=42)),
+    (OtherArgsModel, dict(arg1=42, arg2='abc')),
+])
+def test_single_config_models_fail(tmpdir, cls, config):
+    """ Test fail on passing unsupported config type. """
+    with pytest.raises(ValueError):
+        _ = cls(**config)


Code review comments:
diff_hunk: @@ -51,7 +51,7 @@ def __init__(self,
                  **kwargs) -> object:
         # init superclass
         super().__init__()
-        self.auto_collect_arguments()
+        self._auto_collect_arguments()
path: tests/base/model_template.py
text: @user1:
use 
```
self.save_hyperparameters()
```