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>

pytorch_lightning/loggers/base.py

@@ -1,5 +1,6 @@
 import argparse
 from abc import ABC, abstractmethod
+from argparse import Namespace
 from functools import wraps
 from typing import Union, Optional, Dict, Iterable, Any, Callable, List
 
@@ -41,6 +42,12 @@ class LightningLoggerBase(ABC):
         """
         pass
 
+    def _convert_params(self, params: Union[Dict[str, Any], Namespace]) -> Dict[str, Any]:
+        # in case converting from namespace
+        if isinstance(params, Namespace):
+            params = vars(params)
+        return params
+
     @abstractmethod
     def log_hyperparams(self, params: argparse.Namespace):
         """Record hyperparameters.
@@ -50,11 +57,11 @@ class LightningLoggerBase(ABC):
         """
         pass
 
-    def save(self):
+    def save(self) -> None:
         """Save log data."""
         pass
 
-    def finalize(self, status: str):
+    def finalize(self, status: str) -> None:
         """Do any processing that is necessary to finalize an experiment.
 
         Args:
@@ -62,7 +69,7 @@ class LightningLoggerBase(ABC):
         """
         pass
 
-    def close(self):
+    def close(self) -> None:
         """Do any cleanup that is necessary to close an experiment."""
         pass
 
@@ -72,7 +79,7 @@ class LightningLoggerBase(ABC):
         return self._rank
 
     @rank.setter
-    def rank(self, value: int):
+    def rank(self, value: int) -> None:
         """Set the process rank."""
         self._rank = value
 
@@ -107,23 +114,23 @@ class LoggerCollection(LightningLoggerBase):
     def experiment(self) -> List[Any]:
         return [logger.experiment for logger in self._logger_iterable]
 
-    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
         [logger.log_metrics(metrics, step) for logger in self._logger_iterable]
 
-    def log_hyperparams(self, params: argparse.Namespace):
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         [logger.log_hyperparams(params) for logger in self._logger_iterable]
 
-    def save(self):
+    def save(self) -> None:
         [logger.save() for logger in self._logger_iterable]
 
-    def finalize(self, status: str):
+    def finalize(self, status: str) -> None:
         [logger.finalize(status) for logger in self._logger_iterable]
 
-    def close(self):
+    def close(self) -> None:
         [logger.close() for logger in self._logger_iterable]
 
     @LightningLoggerBase.rank.setter
-    def rank(self, value: int):
+    def rank(self, value: int) -> None:
         self._rank = value
         for logger in self._logger_iterable:
             logger.rank = value

pytorch_lightning/loggers/comet.py

@@ -5,9 +5,9 @@ r"""
 CometLogger
 -------------
 """
-import argparse
+from argparse import Namespace
 from logging import getLogger
-from typing import Optional, Dict, Union
+from typing import Optional, Dict, Union, Any
 
 try:
     from comet_ml import Experiment as CometExperiment
@@ -162,15 +162,16 @@ class CometLogger(LightningLoggerBase):
         return self._experiment
 
     @rank_zero_only
-    def log_hyperparams(self, params: argparse.Namespace):
-        self.experiment.log_parameters(vars(params))
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
+        params = self._convert_params(params)
+        self.experiment.log_parameters(params)
 
     @rank_zero_only
     def log_metrics(
             self,
             metrics: Dict[str, Union[torch.Tensor, float]],
             step: Optional[int] = None
-    ):
+    ) -> None:
         # Comet.ml expects metrics to be a dictionary of detached tensors on CPU
         for key, val in metrics.items():
             if is_tensor(val):
@@ -182,7 +183,7 @@ class CometLogger(LightningLoggerBase):
         self._experiment = None
 
     @rank_zero_only
-    def finalize(self, status: str):
+    def finalize(self, status: str) -> None:
         r"""
         When calling self.experiment.end(), that experiment won't log any more data to Comet. That's why, if you need
         to log any more data you need to create an ExistingCometExperiment. For example, to log data when testing your
@@ -199,7 +200,7 @@ class CometLogger(LightningLoggerBase):
         return self.experiment.project_name
 
     @name.setter
-    def name(self, value: str):
+    def name(self, value: str) -> None:
         self.experiment.set_name(value)
 
     @property

pytorch_lightning/loggers/mlflow.py

@@ -23,10 +23,10 @@ Use the logger anywhere in you LightningModule as follows:
         self.logger.experiment.whatever_ml_flow_supports(...)
 
 """
-import argparse
+from argparse import Namespace
 from logging import getLogger
 from time import time
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, Union
 
 try:
     import mlflow
@@ -88,12 +88,13 @@ class MLFlowLogger(LightningLoggerBase):
         return self._run_id
 
     @rank_zero_only
-    def log_hyperparams(self, params: argparse.Namespace):
-        for k, v in vars(params).items():
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
+        params = self._convert_params(params)
+        for k, v in params.items():
             self.experiment.log_param(self.run_id, k, v)
 
     @rank_zero_only
-    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
         timestamp_ms = int(time() * 1000)
         for k, v in metrics.items():
             if isinstance(v, str):
@@ -105,7 +106,7 @@ class MLFlowLogger(LightningLoggerBase):
         pass
 
     @rank_zero_only
-    def finalize(self, status: str = 'FINISHED'):
+    def finalize(self, status: str = 'FINISHED') -> None:
         if status == 'success':
             status = 'FINISHED'
         self.experiment.set_terminated(self.run_id, status)

pytorch_lightning/loggers/neptune.py

@@ -6,7 +6,7 @@ Log using `neptune-logger <https://www.neptune.ml>`_
 NeptuneLogger
 --------------
 """
-import argparse
+from argparse import Namespace
 from logging import getLogger
 from typing import Optional, List, Dict, Any, Union, Iterable
 
@@ -164,8 +164,9 @@ class NeptuneLogger(LightningLoggerBase):
         return self._experiment
 
     @rank_zero_only
-    def log_hyperparams(self, params: argparse.Namespace):
-        for key, val in vars(params).items():
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
+        params = self._convert_params(params)
+        for key, val in params.items():
             self.experiment.set_property(f'param__{key}', val)
 
     @rank_zero_only
@@ -173,7 +174,7 @@ class NeptuneLogger(LightningLoggerBase):
             self,
             metrics: Dict[str, Union[torch.Tensor, float]],
             step: Optional[int] = None
-    ):
+    ) -> None:
         """Log metrics (numeric values) in Neptune experiments
 
         Args:
@@ -184,7 +185,7 @@ class NeptuneLogger(LightningLoggerBase):
             self.log_metric(key, val, step=step)
 
     @rank_zero_only
-    def finalize(self, status: str):
+    def finalize(self, status: str) -> None:
         self.experiment.stop()
 
     @property
@@ -207,7 +208,7 @@ class NeptuneLogger(LightningLoggerBase):
             metric_name: str,
             metric_value: Union[torch.Tensor, float, str],
             step: Optional[int] = None
-    ):
+    ) -> None:
         """Log metrics (numeric values) in Neptune experiments
 
         Args:
@@ -224,7 +225,7 @@ class NeptuneLogger(LightningLoggerBase):
             self.experiment.log_metric(metric_name, x=step, y=metric_value)
 
     @rank_zero_only
-    def log_text(self, log_name: str, text: str, step: Optional[int] = None):
+    def log_text(self, log_name: str, text: str, step: Optional[int] = None) -> None:
         """Log text data in Neptune experiment
 
         Args:
@@ -235,7 +236,7 @@ class NeptuneLogger(LightningLoggerBase):
         self.log_metric(log_name, text, step=step)
 
     @rank_zero_only
-    def log_image(self, log_name: str, image: Union[str, Any], step: Optional[int] = None):
+    def log_image(self, log_name: str, image: Union[str, Any], step: Optional[int] = None) -> None:
         """Log image data in Neptune experiment
 
         Args:
@@ -250,7 +251,7 @@ class NeptuneLogger(LightningLoggerBase):
             self.experiment.log_image(log_name, x=step, y=image)
 
     @rank_zero_only
-    def log_artifact(self, artifact: str, destination: Optional[str] = None):
+    def log_artifact(self, artifact: str, destination: Optional[str] = None) -> None:
         """Save an artifact (file) in Neptune experiment storage.
 
         Args:
@@ -261,7 +262,7 @@ class NeptuneLogger(LightningLoggerBase):
         self.experiment.log_artifact(artifact, destination)
 
     @rank_zero_only
-    def set_property(self, key: str, value: Any):
+    def set_property(self, key: str, value: Any) -> None:
         """Set key-value pair as Neptune experiment property.
 
         Args:
@@ -271,7 +272,7 @@ class NeptuneLogger(LightningLoggerBase):
         self.experiment.set_property(key, value)
 
     @rank_zero_only
-    def append_tags(self, tags: Union[str, Iterable[str]]):
+    def append_tags(self, tags: Union[str, Iterable[str]]) -> None:
         """appends tags to neptune experiment
 
         Args:

pytorch_lightning/loggers/tensorboard.py

@@ -1,8 +1,7 @@
-import argparse
 import csv
 import os
 from argparse import Namespace
-from typing import Optional, Dict, Union
+from typing import Optional, Dict, Union, Any
 from warnings import warn
 
 import torch
@@ -100,14 +99,8 @@ class TensorBoardLogger(LightningLoggerBase):
         return self._experiment
 
     @rank_zero_only
-    def log_hyperparams(self, params: argparse.Namespace):
-        if params is None:
-            return
-
-        # in case converting from namespace
-        if isinstance(params, Namespace):
-            params = vars(params)
-        params = dict(params)
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
+        params = self._convert_params(params)
 
         if parse_version(torch.__version__) < parse_version("1.3.0"):
             warn(
@@ -126,14 +119,14 @@ class TensorBoardLogger(LightningLoggerBase):
         self.tags.update(params)
 
     @rank_zero_only
-    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
         for k, v in metrics.items():
             if isinstance(v, torch.Tensor):
                 v = v.item()
             self.experiment.add_scalar(k, v, step)
 
     @rank_zero_only
-    def save(self):
+    def save(self) -> None:
         try:
             self.experiment.flush()
         except AttributeError:
@@ -156,7 +149,7 @@ class TensorBoardLogger(LightningLoggerBase):
                 writer.writerow({'key': k, 'value': v})
 
     @rank_zero_only
-    def finalize(self, status: str):
+    def finalize(self, status: str) -> None:
         self.save()
 
     @property

pytorch_lightning/loggers/test_tube.py

@@ -1,5 +1,5 @@
-import argparse
-from typing import Optional, Dict, Any
+from argparse import Namespace
+from typing import Optional, Dict, Any, Union
 
 try:
     from test_tube import Experiment
@@ -92,32 +92,33 @@ class TestTubeLogger(LightningLoggerBase):
         return self._experiment
 
     @rank_zero_only
-    def log_hyperparams(self, params: argparse.Namespace):
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         # TODO: HACK figure out where this is being set to true
         self.experiment.debug = self.debug
-        self.experiment.argparse(params)
+        params = self._convert_params(params)
+        self.experiment.argparse(Namespace(**params))
 
     @rank_zero_only
-    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
         # TODO: HACK figure out where this is being set to true
         self.experiment.debug = self.debug
         self.experiment.log(metrics, global_step=step)
 
     @rank_zero_only
-    def save(self):
+    def save(self) -> None:
         # TODO: HACK figure out where this is being set to true
         self.experiment.debug = self.debug
         self.experiment.save()
 
     @rank_zero_only
-    def finalize(self, status: str):
+    def finalize(self, status: str) -> None:
         # TODO: HACK figure out where this is being set to true
         self.experiment.debug = self.debug
         self.save()
         self.close()
 
     @rank_zero_only
-    def close(self):
+    def close(self) -> None:
         # TODO: HACK figure out where this is being set to true
         self.experiment.debug = self.debug
         if not self.debug:
@@ -129,7 +130,7 @@ class TestTubeLogger(LightningLoggerBase):
         return self._rank
 
     @rank.setter
-    def rank(self, value: int):
+    def rank(self, value: int) -> None:
         self._rank = value
         if self._experiment is not None:
             self.experiment.rank = value

pytorch_lightning/loggers/wandb.py

@@ -5,9 +5,9 @@ r"""
 WandbLogger
 -------------
 """
-import argparse
 import os
-from typing import Optional, List, Dict
+from argparse import Namespace
+from typing import Optional, List, Dict, Union, Any
 
 import torch.nn as nn
 
@@ -91,17 +91,18 @@ class WandbLogger(LightningLoggerBase):
         wandb.watch(model, log=log, log_freq=log_freq)
 
     @rank_zero_only
-    def log_hyperparams(self, params: argparse.Namespace):
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
+        params = self._convert_params(params)
         self.experiment.config.update(params)
 
     @rank_zero_only
-    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
         if step is not None:
             metrics['global_step'] = step
         self.experiment.log(metrics)
 
     @rank_zero_only
-    def finalize(self, status: str = 'success'):
+    def finalize(self, status: str = 'success') -> None:
         try:
             exit_code = 0 if status == 'success' else 1
             wandb.join(exit_code)