Lightning-AI / Lightning-AI/pytorch-lightning
Log `TensorBoard` histograms
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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