Fix TensorBoardLogger's validation of example input when logging graph (#15323)

Author: awaelchli

src/pytorch_lightning/CHANGELOG.md

@@ -46,7 +46,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
--
+- Fixed `TensorBoardLogger` not validating the input array type when logging the model graph ([#15323](https://github.com/Lightning-AI/lightning/pull/15323))
 
 -
 

src/pytorch_lightning/core/module.py

@@ -109,7 +109,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.precision: Union[int, str] = 32
 
         # optionally can be set by user
-        self._example_input_array = None
+        self._example_input_array: Optional[Union[Tensor, Tuple, Dict]] = None
         self._current_fx_name: Optional[str] = None
         self._automatic_optimization: bool = True
         self._truncated_bptt_steps: int = 0
@@ -189,7 +189,7 @@ def trainer(self, trainer: Optional["pl.Trainer"]) -> None:
         self._trainer = trainer
 
     @property
-    def example_input_array(self) -> Any:
+    def example_input_array(self) -> Optional[Union[Tensor, Tuple, Dict]]:
         """The example input array is a specification of what the module can consume in the :meth:`forward` method.
         The return type is interpreted as follows:
 
@@ -203,7 +203,7 @@ def example_input_array(self) -> Any:
         return self._example_input_array
 
     @example_input_array.setter
-    def example_input_array(self, example: Any) -> None:
+    def example_input_array(self, example: Optional[Union[Tensor, Tuple, Dict]]) -> None:
         self._example_input_array = example
 
     @property

src/pytorch_lightning/loggers/tensorboard.py

@@ -235,21 +235,27 @@ def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None)
 
     @rank_zero_only
     def log_graph(self, model: "pl.LightningModule", input_array: Optional[Tensor] = None) -> None:
-        if self._log_graph:
-            if input_array is None:
-                input_array = model.example_input_array
-
-            if input_array is not None:
-                input_array = model._on_before_batch_transfer(input_array)
-                input_array = model._apply_batch_transfer_handler(input_array)
-                with pl.core.module._jit_is_scripting():
-                    self.experiment.add_graph(model, input_array)
-            else:
-                rank_zero_warn(
-                    "Could not log computational graph since the"
-                    " `model.example_input_array` attribute is not set"
-                    " or `input_array` was not given",
-                )
+        if not self._log_graph:
+            return
+
+        input_array = model.example_input_array if input_array is None else input_array
+
+        if input_array is None:
+            rank_zero_warn(
+                "Could not log computational graph to TensorBoard: The `model.example_input_array` attribute"
+                " is not set or `input_array` was not given."
+            )
+        elif not isinstance(input_array, (Tensor, tuple)):
+            rank_zero_warn(
+                "Could not log computational graph to TensorBoard: The `input_array` or `model.example_input_array`"
+                f" has type {type(input_array)} which can't be traced by TensorBoard. Make the input array a tuple"
+                f" representing the positional arguments to the model's `forward()` implementation."
+            )
+        else:
+            input_array = model._on_before_batch_transfer(input_array)
+            input_array = model._apply_batch_transfer_handler(input_array)
+            with pl.core.module._jit_is_scripting():
+                self.experiment.add_graph(model, input_array)
 
     @rank_zero_only
     def save(self) -> None:

tests/tests_pytorch/loggers/test_tensorboard.py

@@ -227,8 +227,13 @@ def test_tensorboard_log_graph_warning_no_example_input_array(tmpdir):
     logger = TensorBoardLogger(tmpdir, log_graph=True)
     with pytest.warns(
         UserWarning,
-        match="Could not log computational graph since the `model.example_input_array`"
-        " attribute is not set or `input_array` was not given",
+        match="Could not log computational graph to TensorBoard: The `model.example_input_array` .* was not given",
+    ):
+        logger.log_graph(model)
+
+    model.example_input_array = dict(x=1, y=2)
+    with pytest.warns(
+        UserWarning, match="Could not log computational graph to TensorBoard: .* can't be traced by TensorBoard"
     ):
         logger.log_graph(model)