Log TensorBoard histograms #19743
Description
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.
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!)
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!