diff --git a/CHANGELOG.md b/CHANGELOG.md index 9e7655f0a..3e190326f 100644 --- a/CHANGELOG.md +++ b/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`. diff --git a/src/lightly_train/_callbacks/tqdm_progress_bar.py b/src/lightly_train/_callbacks/tqdm_progress_bar.py index be27519e0..13139fb81 100644 --- a/src/lightly_train/_callbacks/tqdm_progress_bar.py +++ b/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 diff --git a/src/lightly_train/_methods/method.py b/src/lightly_train/_methods/method.py index 950681549..b60754654 100644 --- a/src/lightly_train/_methods/method.py +++ b/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