Add step parameter to TensorBoardLogger.log_hyperparams (#20176)

* Add step to TensorBoardLogger.log_hyperparams

The metrics that get logged using this method are always logged to step 0, so no meaningful graph is made

* Change expected step value to None

---------

Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Luca Antiga <[email protected]>

Author: 3 people

src/lightning/fabric/loggers/tensorboard.py

@@ -220,15 +220,19 @@ def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None)
     @override
     @rank_zero_only
     def log_hyperparams(
-        self, params: Union[dict[str, Any], Namespace], metrics: Optional[dict[str, Any]] = None
+        self,
+        params: Union[dict[str, Any], Namespace],
+        metrics: Optional[dict[str, Any]] = None,
+        step: Optional[int] = None,
     ) -> None:
         """Record hyperparameters. TensorBoard logs with and without saved hyperparameters are incompatible, the
         hyperparameters are then not displayed in the TensorBoard. Please delete or move the previously saved logs to
         display the new ones with hyperparameters.
 
         Args:
-            params: a dictionary-like container with the hyperparameters
+            params: A dictionary-like container with the hyperparameters
             metrics: Dictionary with metric names as keys and measured quantities as values
+            step: Optional global step number for the logged metrics
 
         """
         params = _convert_params(params)
@@ -244,7 +248,7 @@ def log_hyperparams(
             metrics = {"hp_metric": metrics}
 
         if metrics:
-            self.log_metrics(metrics, 0)
+            self.log_metrics(metrics, step)
 
             if _TENSORBOARD_AVAILABLE:
                 from torch.utils.tensorboard.summary import hparams
@@ -253,9 +257,9 @@ def log_hyperparams(
 
             exp, ssi, sei = hparams(params, metrics)
             writer = self.experiment._get_file_writer()
-            writer.add_summary(exp)
-            writer.add_summary(ssi)
-            writer.add_summary(sei)
+            writer.add_summary(exp, step)
+            writer.add_summary(ssi, step)
+            writer.add_summary(sei, step)
 
     @override
     @rank_zero_only

src/lightning/pytorch/loggers/tensorboard.py

@@ -153,15 +153,19 @@ def save_dir(self) -> str:
     @override
     @rank_zero_only
     def log_hyperparams(
-        self, params: Union[dict[str, Any], Namespace], metrics: Optional[dict[str, Any]] = None
+        self,
+        params: Union[dict[str, Any], Namespace],
+        metrics: Optional[dict[str, Any]] = None,
+        step: Optional[int] = None,
     ) -> None:
         """Record hyperparameters. TensorBoard logs with and without saved hyperparameters are incompatible, the
         hyperparameters are then not displayed in the TensorBoard. Please delete or move the previously saved logs to
         display the new ones with hyperparameters.
 
         Args:
-            params: a dictionary-like container with the hyperparameters
+            params: A dictionary-like container with the hyperparameters
             metrics: Dictionary with metric names as keys and measured quantities as values
+            step: Optional global step number for the logged metrics
 
         """
         if _OMEGACONF_AVAILABLE:
@@ -175,7 +179,7 @@ def log_hyperparams(
         else:
             self.hparams.update(params)
 
-        return super().log_hyperparams(params=params, metrics=metrics)
+        return super().log_hyperparams(params=params, metrics=metrics, step=step)
 
     @override
     @rank_zero_only

tests/tests_pytorch/trainer/logging_/test_eval_loop_logging.py

@@ -544,7 +544,7 @@ def test_step(self, batch, batch_idx):
         "valid_loss_1",
     }
     assert mock_log_metrics.mock_calls == [
-        call({"hp_metric": -1}, 0),
+        call({"hp_metric": -1}, None),
         call(metrics={"train_loss": ANY, "epoch": 0}, step=0),
         call(metrics={"valid_loss_0_step": ANY, "valid_loss_2": ANY}, step=0),
         call(metrics={"valid_loss_0_step": ANY, "valid_loss_2": ANY}, step=1),