TensorBoardLogger sub_dir parameter for grouping logs (#6195)

* fixed a small typo

* cleaning up

* added sub_dir argument to tensorboard and wrote test

* sub dir arg exclusively for tensorboard, linted

* resolving merge conflict

* resolved merge conflict

* resolved merge conflict

* resolved merge conflict

* resolve merge conflict before revert

* resolving merge conflict

* reverted to pre-lint

* added tensorboard sub_dir test

* pep8 formatting

* removed sub_dir arg from test_all function:

* updated feature description

* typo in doc description

* updated CHANGELOG

* Update pytorch_lightning/loggers/tensorboard.py

Co-authored-by: Justus Schock <12886177+justusschock@users.noreply.github.com>

* swapped argument position

* added expandvars tests

* added expandvars

* removed model init

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix tests

* fix failed test

* Revert "fix failed test"

This reverts commit 50b34c6.

* add env var to test

* fix typo in tests

* fix tests

* for test consistency

* fix typo

* fix typo 2

Co-authored-by: Ubuntu <azureuser@devhenrik.evuifrmjd4lepbj4relcwwu5va.ax.internal.cloudapp.net>
Co-authored-by: Justus Schock <12886177+justusschock@users.noreply.github.com>
Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roger Shieh <sh.rog@protonmail.ch>

CHANGELOG.md

@@ -20,6 +20,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `clip_grad_by_value` support for TPUs ([#7025](https://github.com/PyTorchLightning/pytorch-lightning/pull/7025))
 
 
+- Added `sub_dir` parameter to `TensorBoardLogger` ([#6195](https://github.com/PyTorchLightning/pytorch-lightning/pull/6195))
+
+
 - Added correct `dataloader_idx` to batch transfer hooks ([#6241](https://github.com/PyTorchLightning/pytorch-lightning/pull/6241))
 
 

pytorch_lightning/loggers/tensorboard.py

@@ -62,6 +62,9 @@ class TensorBoardLogger(LightningLoggerBase):
             directory for existing versions, then automatically assigns the next available version.
             If it is a string then it is used as the run-specific subdirectory name,
             otherwise ``'version_${version}'`` is used.
+        sub_dir: Sub-directory to group TensorBoard logs. If a sub_dir argument is passed
+            then logs are saved in ``/save_dir/version/sub_dir/``. Defaults to ``None`` in which
+            logs are saved in ``/save_dir/version/``.
         log_graph: Adds the computational graph to tensorboard. This requires that
             the user has defined the `self.example_input_array` attribute in their
             model.
@@ -83,12 +86,14 @@ def __init__(
         log_graph: bool = False,
         default_hp_metric: bool = True,
         prefix: str = '',
+        sub_dir: Optional[str] = None,
         **kwargs
     ):
         super().__init__()
         self._save_dir = save_dir
         self._name = name or ''
         self._version = version
+        self._sub_dir = sub_dir
         self._log_graph = log_graph
         self._default_hp_metric = default_hp_metric
         self._prefix = prefix
@@ -120,12 +125,20 @@ def log_dir(self) -> str:
         # create a pseudo standard path ala test-tube
         version = self.version if isinstance(self.version, str) else f"version_{self.version}"
         log_dir = os.path.join(self.root_dir, version)
+        if isinstance(self.sub_dir, str):
+            log_dir = os.path.join(log_dir, self.sub_dir)
+        log_dir = os.path.expandvars(log_dir)
+        log_dir = os.path.expanduser(log_dir)
         return log_dir
 
     @property
     def save_dir(self) -> Optional[str]:
         return self._save_dir
 
+    @property
+    def sub_dir(self) -> Optional[str]:
+        return self._sub_dir
+
     @property
     @rank_zero_experiment
     def experiment(self) -> SummaryWriter:

tests/loggers/test_all.py

@@ -48,8 +48,8 @@ def _get_logger_args(logger_class, save_dir):
     return logger_args
 
 
-def _instantiate_logger(logger_class, save_idr, **override_kwargs):
-    args = _get_logger_args(logger_class, save_idr)
+def _instantiate_logger(logger_class, save_dir, **override_kwargs):
+    args = _get_logger_args(logger_class, save_dir)
     args.update(**override_kwargs)
     logger = logger_class(**args)
     return logger
@@ -368,38 +368,38 @@ def test_logger_with_prefix_all(tmpdir, monkeypatch):
     with mock.patch('pytorch_lightning.loggers.comet.comet_ml'), \
          mock.patch('pytorch_lightning.loggers.comet.CometOfflineExperiment'):
         _patch_comet_atexit(monkeypatch)
-        logger = _instantiate_logger(CometLogger, save_idr=tmpdir, prefix=prefix)
+        logger = _instantiate_logger(CometLogger, save_dir=tmpdir, prefix=prefix)
         logger.log_metrics({"test": 1.0}, step=0)
         logger.experiment.log_metrics.assert_called_once_with({"tmp-test": 1.0}, epoch=None, step=0)
 
     # MLflow
     with mock.patch('pytorch_lightning.loggers.mlflow.mlflow'), \
          mock.patch('pytorch_lightning.loggers.mlflow.MlflowClient'):
-        logger = _instantiate_logger(MLFlowLogger, save_idr=tmpdir, prefix=prefix)
+        logger = _instantiate_logger(MLFlowLogger, save_dir=tmpdir, prefix=prefix)
         logger.log_metrics({"test": 1.0}, step=0)
         logger.experiment.log_metric.assert_called_once_with(ANY, "tmp-test", 1.0, ANY, 0)
 
     # Neptune
     with mock.patch('pytorch_lightning.loggers.neptune.neptune'):
-        logger = _instantiate_logger(NeptuneLogger, save_idr=tmpdir, prefix=prefix)
+        logger = _instantiate_logger(NeptuneLogger, save_dir=tmpdir, prefix=prefix)
         logger.log_metrics({"test": 1.0}, step=0)
         logger.experiment.log_metric.assert_called_once_with("tmp-test", 1.0)
 
     # TensorBoard
     with mock.patch('pytorch_lightning.loggers.tensorboard.SummaryWriter'):
-        logger = _instantiate_logger(TensorBoardLogger, save_idr=tmpdir, prefix=prefix)
+        logger = _instantiate_logger(TensorBoardLogger, save_dir=tmpdir, prefix=prefix)
         logger.log_metrics({"test": 1.0}, step=0)
         logger.experiment.add_scalar.assert_called_once_with("tmp-test", 1.0, 0)
 
     # TestTube
     with mock.patch('pytorch_lightning.loggers.test_tube.Experiment'):
-        logger = _instantiate_logger(TestTubeLogger, save_idr=tmpdir, prefix=prefix)
+        logger = _instantiate_logger(TestTubeLogger, save_dir=tmpdir, prefix=prefix)
         logger.log_metrics({"test": 1.0}, step=0)
         logger.experiment.log.assert_called_once_with({"tmp-test": 1.0}, global_step=0)
 
     # WandB
     with mock.patch('pytorch_lightning.loggers.wandb.wandb') as wandb:
-        logger = _instantiate_logger(WandbLogger, save_idr=tmpdir, prefix=prefix)
+        logger = _instantiate_logger(WandbLogger, save_dir=tmpdir, prefix=prefix)
         wandb.run = None
         wandb.init().step = 0
         logger.log_metrics({"test": 1.0}, step=0)

tests/loggers/test_tensorboard.py

@@ -119,6 +119,51 @@ def test_tensorboard_no_name(tmpdir, name):
     assert os.listdir(tmpdir / "version_0")
 
 
+def test_tensorboard_log_sub_dir(tmpdir):
+
+    class TestLogger(TensorBoardLogger):
+        # for reproducibility
+        @property
+        def version(self):
+            return "version"
+
+        @property
+        def name(self):
+            return "name"
+
+    trainer_args = dict(
+        default_root_dir=tmpdir,
+        max_steps=1,
+    )
+
+    # no sub_dir specified
+    save_dir = tmpdir / "logs"
+    logger = TestLogger(save_dir)
+    trainer = Trainer(**trainer_args, logger=logger)
+    assert trainer.logger.log_dir == os.path.join(save_dir, "name", "version")
+
+    # sub_dir specified
+    logger = TestLogger(save_dir, sub_dir="sub_dir")
+    trainer = Trainer(**trainer_args, logger=logger)
+    assert trainer.logger.log_dir == os.path.join(save_dir, "name", "version", "sub_dir")
+
+    # test home dir (`~`) handling
+    save_dir = "~/tmp"
+    explicit_save_dir = os.path.expanduser(save_dir)
+    logger = TestLogger(save_dir, sub_dir="sub_dir")
+    trainer = Trainer(**trainer_args, logger=logger)
+    assert trainer.logger.log_dir == os.path.join(explicit_save_dir, "name", "version", "sub_dir")
+
+    # test env var (`$`) handling
+    test_env_dir = "some_directory"
+    os.environ["test_env_dir"] = test_env_dir
+    save_dir = "$test_env_dir/tmp"
+    explicit_save_dir = f"{test_env_dir}/tmp"
+    logger = TestLogger(save_dir, sub_dir="sub_dir")
+    trainer = Trainer(**trainer_args, logger=logger)
+    assert trainer.logger.log_dir == os.path.join(explicit_save_dir, "name", "version", "sub_dir")
+
+
 @pytest.mark.parametrize("step_idx", [10, None])
 def test_tensorboard_log_metrics(tmpdir, step_idx):
     logger = TensorBoardLogger(tmpdir)