Group the metrics in Tensorboard

[mazurkin]

Bug description

I cannot understand what needs to be done to group the graphs for the same metric:

        self.log_dict(
            dictionary={'losses/a': 0.1, 'losses/b': 0.2},
            batch_size=1,
            add_dataloader_idx=False,
            prog_bar=False,
            on_step=True,
            on_epoch=False,
            sync_dist=False,
            rank_zero_only=True,
        )

        self.log(
            name='losses/c',
            value=0.3,
            batch_size=1,
            add_dataloader_idx=False,
            prog_bar=False,
            on_step=True,
            on_epoch=False,
            sync_dist=False,
            rank_zero_only=True,
        )

        self.log(
            name='losses/d',
            value=0.3,
            batch_size=1,
            add_dataloader_idx=False,
            prog_bar=False,
            on_step=True,
            on_epoch=False,
            sync_dist=False,
            rank_zero_only=True,
        )

Result in:

I would really like to see all 4 lines on the same plot.

version: 2.6.1

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

self.log_dict(
            dictionary={'losses/a': 0.1, 'losses/b': 0.2},
            batch_size=1,
            add_dataloader_idx=False,
            prog_bar=False,
            on_step=True,
            on_epoch=False,
            sync_dist=False,
            rank_zero_only=True,
        )

        self.log(
            name='losses/c',
            value=0.3,
            batch_size=1,
            add_dataloader_idx=False,
            prog_bar=False,
            on_step=True,
            on_epoch=False,
            sync_dist=False,
            rank_zero_only=True,
        )

        self.log(
            name='losses/d',
            value=0.3,
            batch_size=1,
            add_dataloader_idx=False,
            prog_bar=False,
            on_step=True,
            on_epoch=False,
            sync_dist=False,
            rank_zero_only=True,
        )

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version: 2.6.1
#- PyTorch Version: 2.9.1
#- Python version: 3.12
#- OS: Linux 
#- CUDA/cuDNN version: 13.0
#- How you installed Lightning: poetry

More info

No response

cc @ethanwharris @lantiga

[mazurkin]

I know I could use SummaryWriter directly and I do it when I do not use Lightning. But here I would like to stay in Lightning's universe and API

https://docs.pytorch.org/docs/stable/tensorboard.html#torch.utils.tensorboard.writer.SummaryWriter.add_scalars

[mazurkin]

this works fine for me, would be nice just to have some standard flag in the constructor for this:

class MultiTbLogger(pl.loggers.TensorBoardLogger):

    ROOT_SPLITTER: t.Final[str] = '/'

    @t.override
    @pl.utilities.rank_zero_only
    def log_metrics(self, metrics: t.Mapping[str, float], step: t.Optional[int] = None) -> None:
        assert pl.utilities.rank_zero_only.rank == 0, 'experiment tried to log from global_rank != 0'

        collection_metrics: dict[str, dict[str, float]] = collections.defaultdict(dict)
        individual_metrics: dict[str, float] = dict()

        for name, value in metrics.items():
            #  split the name into the metric root name and metric value name
            name_parts: list[str] = name.split(self.ROOT_SPLITTER, maxsplit=1)

            match len(name_parts):
                case 2:
                    collection_metrics[name_parts[0]][name_parts[1]] = value
                case 1:
                    individual_metrics[name] = value
                case name_parts_count:
                    raise RuntimeError('the name split failed (empty name?): ' + str(name_parts_count))

        # log all collection metrics
        # the method signature claim it does not support the nested dicts
        # but the implementation actually does (refer to the sources)
        # noinspection PyTypeChecker
        super().log_metrics(collection_metrics, step)

        # log individual metrics
        super().log_metrics(individual_metrics, step)