Lightning-AI / Lightning-AI/pytorch-lightning

Log `TensorBoard` histograms

Open
#19,743 3 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

feature logger: tensorboard
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

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.

- [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 @lantiga @borda

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with lightning/pytorch/core/module.py at __to_tensor and lightning/fabric/loggers/tensorboard.py at the existing scalar logging path. Read SummaryWriter.add_histogram and compare the proposed handling of non-empty one-dimensional tensors with current scalar and dict behavior. Done means supported histogram values pass validation and are emitted through TensorBoard without breaking existing logging.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, observability
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.