Fix misleading data_time / batch_time / data_wait metrics on CUDA (async bleed-through)

CHANGELOG.md

@@ -41,6 +41,7 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Fixed
 
+- Fix `data_time`, `batch_time` and `data_wait` to report correct metrics on cuda.
 - Fix ONNX export for DINOv3 EoMT panoptic, semantic, and instance segmentation models
   by switching to dynamo-based export. Requires torch >= 2.5.0.
 - Fix PicoDet fine-tuning with mismatched `num_classes`.

src/lightly_train/_callbacks/tqdm_progress_bar.py

@@ -7,52 +7,24 @@
 #
 from __future__ import annotations
 
-import time
-from typing import Any
-
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import TQDMProgressBar
 
 
 class DataWaitTQDMProgressBar(TQDMProgressBar):
-    """
-    Customizes the progress bar to include compute efficiency.
-    """
+    """Customizes the progress bar to include compute efficiency."""
 
     def __init__(self) -> None:
         super().__init__(refresh_rate=5)
-        self.batch_start_time: float | None = None
-        self.batch_end_time: float | None = None
-        self.data_time: float | None = None
-        self.batch_time: float | None = None
-
-    def on_train_batch_start(
-        self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int
-    ) -> None:
-        self.batch_start_time = time.perf_counter()
-        if self.batch_end_time is not None:
-            self.data_time = self.batch_start_time - self.batch_end_time
-        return super().on_train_batch_start(trainer, pl_module, batch, batch_idx)
-
-    def on_train_batch_end(
-        self,
-        trainer: Trainer,
-        pl_module: LightningModule,
-        outputs: Any,
-        batch: Any,
-        batch_idx: int,
-    ) -> None:
-        self.batch_end_time = time.perf_counter()
-        if self.batch_start_time is not None:
-            self.batch_time = self.batch_end_time - self.batch_start_time
-        return super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)
 
     def get_metrics(
         self, trainer: Trainer, pl_module: LightningModule
     ) -> dict[str, int | str | float | dict[str, float]]:
         metrics = super().get_metrics(trainer, pl_module)
-        if self.batch_time is not None and self.data_time is not None:
-            if self.batch_time + self.data_time > 0:
-                data_wait = self.data_time / (self.batch_time + self.data_time)
+        batch_time = trainer.callback_metrics.get("profiling/batch_time")
+        data_time = trainer.callback_metrics.get("profiling/data_time")
+        if batch_time is not None and data_time is not None:
+            if batch_time + data_time > 0:
+                data_wait = data_time / (batch_time + data_time)
                 metrics["data_wait"] = f"{data_wait:.1%}"
         return metrics

src/lightly_train/_methods/method.py

@@ -12,6 +12,7 @@
 from dataclasses import dataclass
 from typing import Any, Literal, Mapping
 
+import torch
 from lightly.utils.scheduler import CosineWarmupScheduler
 from pytorch_lightning import LightningModule
 from pytorch_lightning.loggers import WandbLogger as LightningWandbLogger
@@ -44,7 +45,10 @@ class TrainingStepResult:
 @dataclass
 class BatchStartEndTime:
     batch_start_s: float | None = None
-    batch_end_s: float | None = None
+    batch_end_s: float | None = None  # CPU-only
+    # CUDA-only: events bracketing the GPU step.
+    _start_event: torch.cuda.Event | None = None
+    _end_event: torch.cuda.Event | None = None
 
 
 class Method(LightningModule):
@@ -191,27 +195,44 @@ def _log_example_views(self, train_batch: Batch) -> None:
                 )
 
     def _log_time_batch_start(self) -> None:
-        self.batch_start_end_time.batch_start_s = time.perf_counter()
-        if self.batch_start_end_time.batch_end_s is not None:
-            assert (
-                self.batch_start_end_time.batch_start_s
-                > self.batch_start_end_time.batch_end_s
-            )
-            self.log(
-                "profiling/data_time",
-                self.batch_start_end_time.batch_start_s
-                - self.batch_start_end_time.batch_end_s,
-            )
+        now = time.perf_counter()
+        t = self.batch_start_end_time
+
+        if self.device.type == "cuda":
+            # 1. Resolve previous step's GPU duration from events.
+            if (
+                t._start_event is not None
+                and t._end_event is not None
+                and t._end_event.query()  # type: ignore[no-untyped-call]
+            ):
+                batch_time = t._start_event.elapsed_time(t._end_event) / 1e3  # type: ignore[no-untyped-call]
+                self.log("profiling/batch_time", batch_time)
+
+                # 2. data_time = wall gap (start → start) minus GPU duration.
+                if t.batch_start_s is not None:
+                    wall_gap = now - t.batch_start_s
+                    self.log("profiling/data_time", max(0.0, wall_gap - batch_time))
+
+            # 3. Begin new step: record start event.
+            t._start_event = torch.cuda.Event(enable_timing=True)  # type: ignore[no-untyped-call]
+            t._start_event.record()  # type: ignore[no-untyped-call]
+        else:
+            # CPU: data_time is simply the gap between previous step end and now.
+            if t.batch_end_s is not None:
+                self.log("profiling/data_time", now - t.batch_end_s)
+
+        t.batch_start_s = now
 
     def _log_time_batch_end(self) -> None:
-        self.batch_start_end_time.batch_end_s = time.perf_counter()
-        if self.batch_start_end_time.batch_start_s is not None:
-            assert (
-                self.batch_start_end_time.batch_end_s
-                > self.batch_start_end_time.batch_start_s
-            )
-            self.log(
-                "profiling/batch_time",
-                self.batch_start_end_time.batch_end_s
-                - self.batch_start_end_time.batch_start_s,
-            )
+        t = self.batch_start_end_time
+
+        if self.device.type == "cuda":
+            # Record end event — queried at next step start.
+            t._end_event = torch.cuda.Event(enable_timing=True)  # type: ignore[no-untyped-call]
+            t._end_event.record()  # type: ignore[no-untyped-call]
+        else:
+            # CPU: batch_time is just wall-clock step duration.
+            now = time.perf_counter()
+            if t.batch_start_s is not None:
+                self.log("profiling/batch_time", now - t.batch_start_s)
+            t.batch_end_s = now