TensorBoardLogger sub_dir parameter for grouping logs

CHANGELOG.md

@@ -13,6 +13,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))
+
+
 ### Changed
 
 - Log epoch metrics before the `on_evaluation_end` hook ([#7272](https://github.com/PyTorchLightning/pytorch-lightning/pull/7272))

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
@@ -119,13 +124,18 @@ 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)
+        sub_dir = self.sub_dir if isinstance(self.sub_dir, str) else ""
+        log_dir = os.path.join(self.root_dir, version, sub_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,38 @@ 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"
+
+    model = BoringModel()
+    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)
+    trainer.fit(model)
+    assert trainer.logger.log_dir[:-1] == save_dir / "name" / "version"
+
+    # sub_dir specified
+    logger = TestLogger(save_dir, sub_dir="sub_dir")
+    trainer = Trainer(**trainer_args, logger=logger)
+    trainer.fit(model)
+    assert trainer.logger.log_dir == save_dir / "name" / "version" / "sub_dir"
+
+
 @pytest.mark.parametrize("step_idx", [10, None])
 def test_tensorboard_log_metrics(tmpdir, step_idx):
     logger = TensorBoardLogger(tmpdir)