Log multi-value tensors as histograms

Author: dominicgkerr

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

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)