Repository: Lightning-AI/pytorch-lightning
Pull request: TensorBoardLogger sub_dir parameter for grouping logs
## What does this PR do?

**TLDR:** Adds a `sub_dir` parameter for the `TensorBoardLogger` useful for grouping multiple runs in one folder as per #6019.

Although in #6019 I advocated for a general `sub_dir` arg for all logger instances to I eventually added it as a `TensorBoardLogger`-only parameter. Therefore, this PR only solves the "group multiple tensorboard logs into one"-aspect and not the "keep different logger-type logs in different folders"-aspect. It would still be cool to arrive at the solution described in #6019  - yet that will take a bigger overhaul of the logging module (f.e. adding version parameter to all loggers).

## Before submitting
- [x] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [x] Did you make sure to update the documentation with your changes? (if necessary)
- [x] Did you write any new necessary tests? (not for typos and docs)
- [x] Did you verify new and existing tests pass locally with your changes?
- [x] Did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

 - [x] Is this pull request ready for review? (if not, please submit in draft mode)
 - [x] Check that all items from **Before submitting** are resolved
 - [x] Make sure the title is self-explanatory and the description concisely explains the PR
 - [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Support Tensorboard logs dir structure costumizaton
Feature request:
As a user, i would like to have control over checkpoint dir name and structure.

Motivation: 
From @janhenriklambrechts (janhenrik.lambrechts@gmail.com):

My logging version directory of a single pytorch-lighting project with tensorboard logs looks something like this:

```
├── checkpoints
│   ├── epoch=86-max_acc=0.7446.ckpt
│   └── last.ckpt
├── events.out.tfevents.1613393562.lthpc.29645.0
├── hparams.yaml
├── train_acc_classifier_1
│   └── events.out.tfevents.1613393662.lthpc.29645.8
├── train_acc_classifier_2
│   └── events.out.tfevents.1613393662.lthpc.29645.7
├── train_acc_classifier_3
│   └── events.out.tfevents.1613393662.lthpc.29645.6
├── train_acc_classifier_4
│   └── events.out.tfevents.1613393662.lthpc.29645.5
├── val_acc_classifier_1
│   └── events.out.tfevents.1613393563.lthpc.29645.4
├── val_acc_classifier_2
│   └── events.out.tfevents.1613393563.lthpc.29645.3
├── val_acc_classifier_3
│   └── events.out.tfevents.1613393563.lthpc.29645.2
└── val_acc_classifier_4
    └── events.out.tfevents.1613393563.lthpc.29645.1
```
my network has multiple exits and I thus want to log the val_acc (or train_acc) of all these exits in a single graph which is why I need to log these setups as different runs.
As this is very unorganised and I would want to be able to in one scp -r call get all my tensorboard runs to my laptop.

Therefore I was wondering if there is a way in PL to log not directly in my version_0  folder but in version_0/tb_logs
I would want to be able to tell PL to save my logs in a subfolder tensorboard_logs as such that I only have two folders checkpoints  and tensorboard_logs in this version_0  directory.
I would have assumed that the log_dir  parameter of the TensorboardLogger would be this functionality, however in this case the checkpoints and logs are still together but now in a different main folder.

cc @borda @tchaton @justusschock @awaelchli @edward-io @ananthsub @rohitgr7 @kamil-kaczmarek @Raalsky @Blaizzy

Hints from issue discussion:
I propose a ```sub_dir``` argument for ```Logger``` objects. 

The default behaviour for multiple experiments and loggers right now looks something like this:

```
logs
└── default
    └── version_0
        ├── checkpoints
        │   └── epoch=85.ckpt
        ├── comet_log_0.out
        ├── comet_log_1.out
        ├── neptune_log_0.out
        ├── tensorboard_log_0.out
        ├── tensorboard_log_1.out
        └── tensorboard_log_2.out
```

With `sub_dir="tensorboard_logs"` (and `"comet_logs"`, `"neptune_logs"` respectively) it will look something like this

```
logs
└── default
    └── version_0
        ├── checkpoints
        │   └── epoch=85.ckpt
        ├── comet_logs
        │   ├── comet_log_0.out
        │   └── comet_log_1.out
        ├── neptune_logs
        │   └── neptune_log_0.out
        └── tensorboard_logs
            ├── tensorboard_log_0.out
            ├── tensorboard_log_1.out
            └── tensorboard_log_2.out
```

- I think creating a subdirectory for these logs should also be default behaviour. - given that I see no default usecase for not wanting a `sub_dir` but do consider having subfolders expected behaviour with multiple loggers/experiments. 
- Moreover, checkpoints are automatically grouped into a directory as well even if there is only one checkpoint being saved - so it makes sense from a design point as well.
- If there are good reasons against this being the default, we could leave `sub_dir=None` as the default argument instead.

Happy to submit a PR for this, or discussing concerns with this approach 😊

Commit message: Update pytorch_lightning/loggers/tensorboard.py

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

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 06656907ef442..5a0f43de3a5d1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `checkpoint` parameter to callback's `on_save_checkpoint` hook ([#6072](https://github.com/PyTorchLightning/pytorch-lightning/pull/6072))
 
+- Added `sub_dir` parameter to `TensorBoardLogger` ([#6195](https://github.com/PyTorchLightning/pytorch-lightning/pull/6195))
 
 ### Changed
 
diff --git a/pytorch_lightning/loggers/base.py b/pytorch_lightning/loggers/base.py
index 4fdb5e8c437bf..3288db607201f 100644
--- a/pytorch_lightning/loggers/base.py
+++ b/pytorch_lightning/loggers/base.py
@@ -425,7 +425,7 @@ def __getitem__(self, idx):
 
 
 class DummyLogger(LightningLoggerBase):
-    """ Dummy logger for internal use. Is usefull if we want to disable users
+    """ Dummy logger for internal use. Is useful if we want to disable users
         logger for a feature, but still secure that users code can run """
 
     def __init__(self):
diff --git a/pytorch_lightning/loggers/tensorboard.py b/pytorch_lightning/loggers/tensorboard.py
index 0485868fa2ef1..844517acec4ea 100644
--- a/pytorch_lightning/loggers/tensorboard.py
+++ b/pytorch_lightning/loggers/tensorboard.py
@@ -57,6 +57,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.
@@ -75,6 +78,7 @@ def __init__(
         save_dir: str,
         name: Optional[str] = "default",
         version: Optional[Union[int, str]] = None,
+        sub_dir: Optional[str] = None,
         log_graph: bool = False,
         default_hp_metric: bool = True,
         prefix: str = '',
@@ -84,6 +88,7 @@ def __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
@@ -114,13 +119,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:
diff --git a/tests/loggers/test_all.py b/tests/loggers/test_all.py
index c80dddde2774c..80e936de74736 100644
--- a/tests/loggers/test_all.py
+++ b/tests/loggers/test_all.py
@@ -49,8 +49,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
@@ -369,38 +369,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)
diff --git a/tests/loggers/test_tensorboard.py b/tests/loggers/test_tensorboard.py
index e5e3f231d3ac7..d5e5ec0786b05 100644
--- a/tests/loggers/test_tensorboard.py
+++ b/tests/loggers/test_tensorboard.py
@@ -121,6 +121,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)


Code review comments:
diff_hunk: @@ -114,13 +119,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)
path: pytorch_lightning/loggers/tensorboard.py
text: @user1:
Should we check sub_dir doesn't contain `~` or `..` 

@user2:
I think you can just use the expand options for os in every case here.

@author:
@user2 I am not exactly following. You mean adding the `os.path.expand` option to every logger `log_dir` ?

@user3:
```suggestion
        log_dir = os.path.join(self.root_dir, version, sub_dir)
        log_dir = os.path.expandvars(log_dir)
        log_dir = os.path.expanduser(log_dir)
```

I think this is what justus means, would also need tests for `~` and `$` in paths.