Fixed handling on enums in hyperparams save method, Added unit test, Comments implemented

Author: grajat90

CHANGELOG.md

@@ -224,6 +224,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `XLACheckpointIO` plugin ([#9972](https://github.com/PyTorchLightning/pytorch-lightning/pull/9972))
 
 
+- Added `use_omegaconf` argument to `save_hparams_to_yaml` plugin ([#9170](https://github.com/PyTorchLightning/pytorch-lightning/pull/9170))
+
 
 ### Changed
 
@@ -238,6 +240,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
   Old [neptune-client](https://github.com/neptune-ai/neptune-client) API is supported by `NeptuneClient` from [neptune-contrib](https://github.com/neptune-ai/neptune-contrib) repo.
 
 
+- Parsing of `enums` type hyperparameters to be saved in the `haprams.yaml` file by tensorboard and csv loggers has been fixed and made in line with how omegaconf parses it. ([#9170](https://github.com/PyTorchLightning/pytorch-lightning/pull/9170))
+
+
 - Parsing of the `gpus` Trainer argument has changed: `gpus="n"` (str) no longer selects the GPU index n and instead selects the first n devices. ([#8770](https://github.com/PyTorchLightning/pytorch-lightning/pull/8770))
 
 

pytorch_lightning/core/saving.py

@@ -19,6 +19,7 @@
 import os
 from argparse import Namespace
 from copy import deepcopy
+from enum import Enum
 from typing import Any, Callable, Dict, IO, MutableMapping, Optional, Union
 from warnings import warn
 
@@ -381,6 +382,8 @@ def save_hparams_to_yaml(config_yaml, hparams: Union[dict, Namespace]) -> None:
     # drop paramaters which contain some strange datatypes as fsspec
     for k, v in hparams.items():
         try:
+            if isinstance(v, Enum):
+                v = deepcopy(v.name)
             yaml.dump(v)
         except TypeError:
             warn(f"Skipping '{k}' parameter because it is not possible to safely dump to YAML.")

tests/models/test_hparams.py

@@ -17,13 +17,16 @@
 import pickle
 from argparse import Namespace
 from dataclasses import dataclass
+from enum import Enum
 from unittest import mock
 
 import cloudpickle
 import pytest
 import torch
 from fsspec.implementations.local import LocalFileSystem
 from omegaconf import Container, OmegaConf
+from omegaconf.dictconfig import DictConfig
+
 from torch.utils.data import DataLoader
 
 from pytorch_lightning import LightningModule, Trainer
@@ -477,22 +480,36 @@ def test_hparams_pickle_warning(tmpdir):
 
 
 def test_hparams_save_yaml(tmpdir):
+    class Options(str, Enum):
+        option1name = "option1val"
+        option2name = "option2val"
+        option3name = "option3val"
     hparams = dict(
-        batch_size=32, learning_rate=0.001, data_root="./any/path/here", nasted=dict(any_num=123, anystr="abcd")
+        batch_size=32, learning_rate=0.001, data_root="./any/path/here", nested=dict(any_num=123, anystr="abcd"),
+        switch= Options.option3name
     )
     path_yaml = os.path.join(tmpdir, "testing-hparams.yaml")
 
+    def compare(loadedParams, defaultParams: dict):
+        assert isinstance(loadedParams, (dict, DictConfig))
+        assert loadedParams.keys() == defaultParams.keys()
+        for k,v in defaultParams.items():
+            if isinstance(v, Enum):
+                assert v.name == loadedParams[k]
+            else:
+                assert v == loadedParams[k]
+
     save_hparams_to_yaml(path_yaml, hparams)
-    assert load_hparams_from_yaml(path_yaml, use_omegaconf=False) == hparams
+    compare(load_hparams_from_yaml(path_yaml, use_omegaconf=False), hparams)
 
     save_hparams_to_yaml(path_yaml, Namespace(**hparams))
-    assert load_hparams_from_yaml(path_yaml, use_omegaconf=False) == hparams
+    compare(load_hparams_from_yaml(path_yaml, use_omegaconf=False), hparams)
 
     save_hparams_to_yaml(path_yaml, AttributeDict(hparams))
-    assert load_hparams_from_yaml(path_yaml, use_omegaconf=False) == hparams
+    compare(load_hparams_from_yaml(path_yaml, use_omegaconf=False),hparams)
 
     save_hparams_to_yaml(path_yaml, OmegaConf.create(hparams))
-    assert load_hparams_from_yaml(path_yaml) == hparams
+    compare(load_hparams_from_yaml(path_yaml), hparams)
 
 
 class NoArgsSubClassBoringModel(CustomBoringModel):