Improve val step logging (#7351)

* Fix val step logging

* Add a type

* Fix

* Update CHANGELOG.md

Author: ethanwharris

CHANGELOG.md

@@ -141,6 +141,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Ensure accelerator is valid if running interactively ([#5970](https://github.com/PyTorchLightning/pytorch-lightning/pull/5970))
 - Disabled batch transfer in DP mode ([#6098](https://github.com/PyTorchLightning/pytorch-lightning/pull/6098))
 
+- Changed the behaviour when logging evaluation step metrics to no longer append `/epoch_*` to the metric name ([#7351](https://github.com/PyTorchLightning/pytorch-lightning/pull/7351))
+
+
 ### Deprecated
 
 - Deprecated `outputs` in both `LightningModule.on_train_epoch_end` and `Callback.on_train_epoch_end` hooks ([#7339](https://github.com/PyTorchLightning/pytorch-lightning/pull/7339))

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -43,6 +43,8 @@ def __init__(self, trainer, log_gpu_memory: Optional[str] = None):
         self._cached_results = {stage: EpochResultStore(trainer) for stage in RunningStage}
         self._cached_results[None] = EpochResultStore(trainer)
         self._callback_hook_validator = CallbackHookNameValidator()
+        self._val_log_step: int = 0
+        self._test_log_step: int = 0
 
     @property
     def callback_metrics(self) -> Dict:
@@ -201,7 +203,8 @@ def log_metrics(self, metrics, grad_norm_dic, step=None):
         Args:
             metrics (dict): Metric values
             grad_norm_dic (dict): Gradient norms
-            step (int): Step for which metrics should be logged. Default value corresponds to `self.global_step`
+            step (int): Step for which metrics should be logged. Default value is `self.global_step` during training or
+                the total validation / test log step count during validation and testing.
         """
         # add gpu memory
         if self.trainer._device_type == DeviceType.GPU and self.log_gpu_memory:
@@ -371,3 +374,31 @@ def log_train_step_metrics(self, batch_output):
             if len(batch_log_metrics) > 0 or len(grad_norm_dic) > 0:
                 self.log_metrics(batch_log_metrics, grad_norm_dic)
                 self._callback_metrics.update(batch_log_metrics)
+
+    @property
+    def evaluation_log_step(self) -> Optional[int]:
+        if self.trainer.state.stage is RunningStage.VALIDATING:
+            return self._val_log_step
+        elif self.trainer.state.stage is RunningStage.TESTING:
+            return self._test_log_step
+        else:
+            return None
+
+    def increment_evaluation_log_step(self) -> None:
+        if self.trainer.state.stage is RunningStage.VALIDATING:
+            self._val_log_step += 1
+        elif self.trainer.state.stage is RunningStage.TESTING:
+            self._test_log_step += 1
+
+    def log_evaluation_step_metrics(self) -> None:
+        if self.trainer.sanity_checking:
+            return
+        _, batch_log_metrics = self.cached_results.update_logger_connector()
+
+        # logs user requested information to logger
+        if len(batch_log_metrics) > 0:
+            kwargs = dict() if "step" in batch_log_metrics else dict(step=self.evaluation_log_step)
+            self.log_metrics(batch_log_metrics, {}, **kwargs)
+
+        # increment the step even if nothing was logged
+        self.increment_evaluation_log_step()

pytorch_lightning/trainer/evaluation_loop.py

@@ -267,22 +267,3 @@ def on_evaluation_epoch_end(self) -> None:
         self.trainer._cache_logged_metrics()
 
         self.trainer.call_hook('on_epoch_end')
-
-    def log_evaluation_step_metrics(self, batch_idx: int) -> None:
-        if self.trainer.sanity_checking:
-            return
-
-        cached_results = self.trainer.logger_connector.cached_results
-        if cached_results is not None:
-            cached_batch_pbar_metrics, cached_batch_log_metrics = cached_results.update_logger_connector()
-
-            if len(cached_batch_log_metrics) > 0:
-                # make the metrics appear as a different line in the same graph
-                metrics_by_epoch = {}
-                for k, v in cached_batch_log_metrics.items():
-                    metrics_by_epoch[f'{k}/epoch_{self.trainer.current_epoch}'] = v
-
-                self.trainer.logger_connector.log_metrics(metrics_by_epoch, {}, step=batch_idx)
-
-            if len(cached_batch_pbar_metrics) > 0:
-                self.trainer.logger_connector.add_progress_bar_metrics(cached_batch_pbar_metrics)

pytorch_lightning/trainer/trainer.py

@@ -966,7 +966,7 @@ def run_evaluation(self, on_epoch: bool = False) -> _EVALUATE_OUTPUT:
                 self.evaluation_loop.on_evaluation_batch_end(output, batch, batch_idx, dataloader_idx)
 
                 # log batch metrics
-                self.evaluation_loop.log_evaluation_step_metrics(batch_idx)
+                self.logger_connector.log_evaluation_step_metrics()
 
                 # track epoch level outputs
                 dl_outputs = self.track_output_for_epoch_end(dl_outputs, output)

tests/accelerators/test_multi_nodes_gpu.py

@@ -121,8 +121,7 @@ def backward(self, loss, optimizer, optimizer_idx):
         'a2',
         'a_step',
         'a_epoch',
-        'b_step/epoch_0',
-        'b_step/epoch_1',
+        'b_step',
         'b_epoch',
         'epoch',
     }

tests/trainer/logging_/test_eval_loop_logging.py

@@ -76,16 +76,15 @@ def backward(self, loss, optimizer, optimizer_idx):
         'a2',
         'a_step',
         'a_epoch',
-        'b_step/epoch_0',
-        'b_step/epoch_1',
+        'b_step',
         'b_epoch',
         'epoch',
     }
     logged_metrics = set(trainer.logged_metrics.keys())
     assert expected_logged_metrics == logged_metrics
 
     # we don't want to enable val metrics during steps because it is not something that users should do
-    # on purpose DO NOT allow step_b... it's silly to monitor val step metrics
+    # on purpose DO NOT allow b_step... it's silly to monitor val step metrics
     callback_metrics = set(trainer.callback_metrics.keys())
     expected_cb_metrics = {'a', 'a2', 'b', 'a_epoch', 'b_epoch', 'a_step'}
     assert expected_cb_metrics == callback_metrics
@@ -145,8 +144,7 @@ def backward(self, loss, optimizer, optimizer_idx):
         'b_step',
         'b_epoch',
         'c',
-        'd_step/epoch_0',
-        'd_step/epoch_1',
+        'd_step',
         'd_epoch',
         'g',
     }
@@ -294,15 +292,15 @@ def validation_epoch_end(self, outputs) -> None:
 
     # make sure values are correct
     assert trainer.logged_metrics['val_loss_epoch'] == manual_mean
-    assert trainer.callback_metrics['val_loss'] == trainer.logged_metrics['val_loss_step/epoch_0']
+    assert trainer.callback_metrics['val_loss'] == trainer.logged_metrics['val_loss_step']
 
     # make sure correct values were logged
     logged_val = trainer.dev_debugger.logged_metrics
 
     # 3 val batches
-    assert logged_val[0]['val_loss_step/epoch_0'] == model.seen_vals[0]
-    assert logged_val[1]['val_loss_step/epoch_0'] == model.seen_vals[1]
-    assert logged_val[2]['val_loss_step/epoch_0'] == model.seen_vals[2]
+    assert logged_val[0]['val_loss_step'] == model.seen_vals[0]
+    assert logged_val[1]['val_loss_step'] == model.seen_vals[1]
+    assert logged_val[2]['val_loss_step'] == model.seen_vals[2]
 
     # epoch mean
     assert logged_val[3]['val_loss_epoch'] == model.manual_epoch_end_mean
@@ -872,29 +870,29 @@ def get_metrics_at_idx(idx):
         else:
             return mock_calls[idx][2]["metrics"]
 
-    expected = ['valid_loss_0_step/epoch_0', 'valid_loss_2/epoch_0', 'global_step']
+    expected = ['valid_loss_0_step', 'valid_loss_2', 'global_step']
     assert sorted(get_metrics_at_idx(1)) == sorted(expected)
     assert sorted(get_metrics_at_idx(2)) == sorted(expected)
 
     expected = model.val_losses[2]
-    assert get_metrics_at_idx(1)["valid_loss_0_step/epoch_0"] == expected
+    assert get_metrics_at_idx(1)["valid_loss_0_step"] == expected
     expected = model.val_losses[3]
-    assert get_metrics_at_idx(2)["valid_loss_0_step/epoch_0"] == expected
+    assert get_metrics_at_idx(2)["valid_loss_0_step"] == expected
 
     expected = ['valid_loss_0_epoch', 'valid_loss_1', 'epoch', 'global_step']
     assert sorted(get_metrics_at_idx(3)) == sorted(expected)
 
     expected = torch.stack(model.val_losses[2:4]).mean()
     assert get_metrics_at_idx(3)["valid_loss_1"] == expected
-    expected = ['valid_loss_0_step/epoch_1', 'valid_loss_2/epoch_1', 'global_step']
+    expected = ['valid_loss_0_step', 'valid_loss_2', 'global_step']
 
     assert sorted(get_metrics_at_idx(4)) == sorted(expected)
     assert sorted(get_metrics_at_idx(5)) == sorted(expected)
 
     expected = model.val_losses[4]
-    assert get_metrics_at_idx(4)["valid_loss_0_step/epoch_1"] == expected
+    assert get_metrics_at_idx(4)["valid_loss_0_step"] == expected
     expected = model.val_losses[5]
-    assert get_metrics_at_idx(5)["valid_loss_0_step/epoch_1"] == expected
+    assert get_metrics_at_idx(5)["valid_loss_0_step"] == expected
 
     expected = ['valid_loss_0_epoch', 'valid_loss_1', 'epoch', 'global_step']
     assert sorted(get_metrics_at_idx(6)) == sorted(expected)

tests/trainer/logging_/test_train_loop_logging.py

@@ -401,7 +401,7 @@ def val_dataloader(self):
     trainer.fit(model)
 
     generated = set(trainer.logger_connector.logged_metrics)
-    expected = {'a_step', 'a_epoch', 'n_step/epoch_0', 'n_epoch', 'epoch'}
+    expected = {'a_step', 'a_epoch', 'n_step', 'n_epoch', 'epoch'}
 
     assert generated == expected