Description & Motivation

I'd like to be able to log histograms using TensorBoard. Looking at how multiple/single scalars are currently logged by lightning.fabric.loggers.TensorBoardLogger [1], I think SummaryWriter.add_histogram [2] could be called (non-empty torch.Tensor instances) in a similar way

Unfortunately, values passed into LightningModule.log are first type-checked by LightningModule.__to_tensor [3] - this enforces values to be either dict or single-value (0Dim) torch.Tensor instances. So in order to support histograms, this check would need to be loosen slightly...

• [1] https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/loggers/tensorboard.py#L211
• [2] https://pytorch.org/docs/stable/tensorboard.html#torch.utils.tensorboard.writer.SummaryWriter.add_histogram
• [3] https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/core/module.py#L642

Pitch

Firstly, I'm proposing [4] becomes:

if not (torch.numel(value) == 1 or (torch.numel(value) > 0 and value.ndim == 1)):
    raise ValueError(
        f"`self.log({name}, {value})` was called, but the tensor must have a single element,"
        f"or a single non-empty dimension. You can try doing `self.log({name}, {value}.mean())`"
    )

and [1] becomes (say):

try:
    if isinstance(v, dict):
        self.experiment.add_scalars(k, v, step)
    elif isinstance(v, Tensor):
        self.experiment.add_histogram(k, v, step)
    else:
        self.experiment.add_scalar(k, v, step)

# TODO(fabric): specify the possible exception
except Exception as ex:
    raise ValueError(
        f"\n you tried to log {v} which is currently not supported. Try a dict or a scalar/tensor."
    ) from ex

I'm sure you're keen not to modify LightnginModule too much, but I couldn't figure out a way to pass 1+Dim Tensor instances through to the TensorBoardLogger without loosing the type checking.

• [4] https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/core/module.py#L648

Alternatives

As a workaround, I'm currently using a subclass of LightningModule that overloads the .log method [5] with:

def log(self, name: str, value: _METRIC, *args, **kwargs):
    if isinstance(value, Tensor) and value.ndim > 0:
        for logger in self.loggers:
            if isinstance(logger, TensorBoardLogger):
                logger.experiment.add_histogram(name, value, self.global_step)

    else:
        super().log(name, value, *args, **kwargs)

Introducing TensorBoard specifics into the LightningModule class, is a horrible hack (hence wanting to incorporate the "nicer" fix above!)

• [5] https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/core/module.py#L371

Additional context

I'm very happy to open a PR (I have one written up from exploring potential implementations) is the above sounds reasonable?

Many thanks!

cc @Borda