diff --git a/src/lightning/fabric/loggers/tensorboard.py b/src/lightning/fabric/loggers/tensorboard.py index 208244dc38cd3..c159c81854e8e 100644 --- a/src/lightning/fabric/loggers/tensorboard.py +++ b/src/lightning/fabric/loggers/tensorboard.py @@ -203,19 +203,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 diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 8eef1c0e7c6e8..2996c0267de46 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/src/lightning/pytorch/CHANGELOG.md @@ -27,6 +27,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). ### Changed +- Allow `LightningModule` to log `Tensorboard` histograms ([#19851](https://github.com/Lightning-AI/pytorch-lightning/pull/19851)) + + - Default to `RichProgressBar` and `RichModelSummary` if the rich package is available. Fallback to TQDMProgressBar and ModelSummary otherwise ([#20896](https://github.com/Lightning-AI/pytorch-lightning/pull/20896)) diff --git a/src/lightning/pytorch/core/module.py b/src/lightning/pytorch/core/module.py index 85f631ee40f75..1c4aa081a2ca3 100644 --- a/src/lightning/pytorch/core/module.py +++ b/src/lightning/pytorch/core/module.py @@ -661,10 +661,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 diff --git a/tests/tests_pytorch/loggers/test_tensorboard.py b/tests/tests_pytorch/loggers/test_tensorboard.py index 7e02a73c93082..fc347cb560c76 100644 --- a/tests/tests_pytorch/loggers/test_tensorboard.py +++ b/tests/tests_pytorch/loggers/test_tensorboard.py @@ -153,7 +153,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) diff --git a/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py b/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py index 6916eae68e9c0..2120a628e7980 100644 --- a/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py +++ b/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)