Log TensorBoard histograms

src/lightning/fabric/loggers/tensorboard.py

@@ -205,19 +205,22 @@ def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None)
  metrics = _add_prefix(metrics, self._prefix, self.LOGGER_JOIN_CHAR)
 
  for k, v in metrics.items():
- if isinstance(v, Tensor):
+ if isinstance(v, Tensor) and v.ndim == 0:
  v = v.item()
 
- if isinstance(v, dict):
- self.experiment.add_scalars(k, v, step)
- else:
- try:
+ 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
+
+ # 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
 
  @override
  @rank_zero_only

src/lightning/pytorch/CHANGELOG.md

@@ -25,7 +25,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Relaxed the requirement for custom batch samplers to expose `drop_last` for prediction ([#19678](https://github.com/Lightning-AI/pytorch-lightning/pull/19678))
 
--
+- Allow `LightningModule` to log `Tensorboard` histograms ([#19851](https://github.com/Lightning-AI/pytorch-lightning/pull/19851))
 
 ### Deprecated
 

src/lightning/pytorch/core/module.py

@@ -642,10 +642,12 @@ def __to_tensor(self, value: Union[Tensor, numbers.Number], name: str) -> Tensor
  if isinstance(value, Tensor)
  else torch.tensor(value, device=self.device, dtype=_get_default_dtype())
  )
- if not torch.numel(value) == 1:
+
+ # check tensor contains single element (implies value.ndim == 0), or is a non-empty 1D array
+ 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" You can try doing `self.log({name}, {value}.mean())`"
+ 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())`"
  )
  value = value.squeeze()
  return value

tests/tests_pytorch/loggers/test_tensorboard.py

@@ -154,7 +154,13 @@ def name(self):
 @pytest.mark.parametrize("step_idx", [10, None])
 def test_tensorboard_log_metrics(tmp_path, step_idx):
  logger = TensorBoardLogger(tmp_path)
- metrics = {"float": 0.3, "int": 1, "FloatTensor": torch.tensor(0.1), "IntTensor": torch.tensor(1)}
+ metrics = {
+ "float": 0.3,
+ "int": 1,
+ "FloatTensor": torch.tensor(0.1),
+ "IntTensor": torch.tensor(1),
+ "Histogram": torch.tensor([10, 100, 1000]),
+ }
  logger.log_metrics(metrics, step_idx)
 
 

tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py

@@ -640,10 +640,18 @@ def training_step(self, *args):
 
  class TestModel(BoringModel):
  def on_train_start(self):
- self.log("foo", torch.tensor([1.0, 2.0]))
+ self.log("foo", torch.tensor([])) # empty
 
  model = TestModel()
- with pytest.raises(ValueError, match="tensor must have a single element"):
+ with pytest.raises(ValueError, match="tensor must have a single element, or a single non-empty dimension."):
+ trainer.fit(model)
+
+ class TestModel(BoringModel):
+ def on_train_start(self):
+ self.log("foo", torch.tensor([[1.0], [2.0]])) # too-many dimensions
+
+ model = TestModel()
+ with pytest.raises(ValueError, match="tensor must have a single element, or a single non-empty dimension."):
  trainer.fit(model)