feat: optionally preserve sequences in WandB logs

Author: nightcityblade

ignite/handlers/base_logger.py

@@ -116,7 +116,11 @@ def global_step_transform(engine: Engine, event_name: str | Events) -> int:
         self.state_attributes = state_attributes
 
     def _setup_output_metrics_state_attrs(
-        self, engine: Engine, log_text: bool | None = False, key_tuple: bool | None = True
+        self,
+        engine: Engine,
+        log_text: bool | None = False,
+        key_tuple: bool | None = True,
+        flatten_sequences: bool = True,
     ) -> dict[Any, Any]:
         """Helper method to setup metrics and state attributes to log"""
         metrics_state_attrs = OrderedDict()
@@ -144,7 +148,7 @@ def _setup_output_metrics_state_attrs(
         if self.state_attributes is not None:
             metrics_state_attrs.update({name: getattr(engine.state, name, None) for name in self.state_attributes})
 
-        metrics_state_attrs_dict: dict[Any, str | float | numbers.Number] = OrderedDict()
+        metrics_state_attrs_dict: dict[Any, str | float | numbers.Number | list[numbers.Number]] = OrderedDict()
 
         def key_tuple_fn(parent_key: str | tuple[str, ...] | None, *args: str) -> tuple[str, ...]:
             if parent_key is None:
@@ -161,19 +165,25 @@ def key_str_fn(parent_key: str, *args: str) -> str:
 
         def handle_value_fn(
             value: str | int | float | numbers.Number | torch.Tensor,
-        ) -> None | str | float | numbers.Number:
+        ) -> None | str | float | numbers.Number | list[numbers.Number]:
             if isinstance(value, numbers.Number):
                 return value
             elif isinstance(value, torch.Tensor) and value.ndimension() == 0:
                 return value.item()
+            elif not flatten_sequences and isinstance(value, Sequence):
+                scalar_values = [item.item() if isinstance(item, torch.Tensor) else item for item in value]
+                if all(isinstance(item, numbers.Number) for item in scalar_values):
+                    return scalar_values
             else:
                 if isinstance(value, str) and log_text:
                     return value
                 else:
                     warnings.warn(f"Logger output_handler can not log metrics value type {type(value)}")
             return None
 
-        metrics_state_attrs_dict = _flatten_dict(metrics_state_attrs, key_fn, handle_value_fn, parent_key=self.tag)
+        metrics_state_attrs_dict = _flatten_dict(
+            metrics_state_attrs, key_fn, handle_value_fn, parent_key=self.tag, flatten_sequences=flatten_sequences
+        )
         return metrics_state_attrs_dict
 
 
@@ -182,13 +192,14 @@ def _flatten_dict(
     key_fn: Callable,
     value_fn: Callable,
     parent_key: str | tuple[str, ...] | None = None,
+    flatten_sequences: bool = True,
 ) -> dict:
     items = {}
     for key, value in in_dict.items():
         new_key = key_fn(parent_key, key)
         if isinstance(value, Mapping):
-            items.update(_flatten_dict(value, key_fn, value_fn, new_key))
-        elif any(
+            items.update(_flatten_dict(value, key_fn, value_fn, new_key, flatten_sequences))
+        elif flatten_sequences and any(
             [
                 isinstance(value, tuple) and hasattr(value, "_fields"),  # namedtuple
                 not isinstance(value, str) and isinstance(value, Sequence),

ignite/handlers/wandb_logger.py

@@ -174,6 +174,8 @@ class OutputHandler(BaseOutputHandler):
             uses function output as global_step. To setup global step from another engine, please use
             :meth:`~ignite.handlers.wandb_logger.global_step_from_engine`.
         sync: Deprecated, has no function. Argument is kept here for compatibility with existing code.
+        flatten_sequences: Whether to flatten sequences into separate metrics. Set to ``False`` to pass sequences of
+            scalar values directly to Weights & Biases.
 
     Examples:
         .. code-block:: python
@@ -282,8 +284,10 @@ def __init__(
         global_step_transform: Callable[[Engine, str | Events], int] | None = None,
         sync: bool | None = None,
         state_attributes: list[str] | None = None,
+        flatten_sequences: bool = True,
     ):
         super().__init__(tag, metric_names, output_transform, global_step_transform, state_attributes)
+        self.flatten_sequences = flatten_sequences
         if sync is not None:
             warn("The sync argument for the WandBLoggers is no longer used, and may be removed in the future")
 
@@ -297,7 +301,9 @@ def __call__(self, engine: Engine, logger: WandBLogger, event_name: str | Events
                 f"global_step must be int, got {type(global_step)}. Please check the output of global_step_transform."
             )
 
-        metrics = self._setup_output_metrics_state_attrs(engine, log_text=True, key_tuple=False)
+        metrics = self._setup_output_metrics_state_attrs(
+            engine, log_text=True, key_tuple=False, flatten_sequences=self.flatten_sequences
+        )
         logger.log(metrics, step=global_step)
 
 

tests/ignite/handlers/test_wandb_logger.py

@@ -131,6 +131,11 @@ def test_output_handler_metric_names():
     wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
     mock_logger.log.assert_called_once_with({f"tag/a/{i}": v for i, v in enumerate(data)}, step=7)
 
+    wrapper = OutputHandler("tag", metric_names=["a"], flatten_sequences=False)
+    mock_logger.log = MagicMock()
+    wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
+    mock_logger.log.assert_called_once_with({"tag/a": data}, step=7)
+
     wrapper = OutputHandler("tag", metric_names="all")
     mock_engine = MagicMock()
     mock_engine.state = State(