feature: Record batch time

src/refiners/fluxion/utils.py

@@ -229,8 +229,7 @@ def load_tensors(path: Path | str, /, device: Device | str = "cpu") -> dict[str,
         tensors = torch.load(path, map_location=device, weights_only=True)  # type: ignore
 
     assert isinstance(tensors, dict) and all(
-        isinstance(key, str) and isinstance(value, Tensor)
-        for key, value in tensors.items()  # type: ignore
+        isinstance(key, str) and isinstance(value, Tensor) for key, value in tensors.items()  # type: ignore
     ), "Invalid tensor file, expected a dict[str, Tensor]"
 
     return cast(dict[str, Tensor], tensors)

src/refiners/foundationals/latent_diffusion/stable_diffusion_1/unet.py

@@ -165,7 +165,8 @@ def __init__(self, device: Device | str | None = None, dtype: DType | None = Non
 class SD1UNet(fl.Chain):
     """Stable Diffusion 1.5 U-Net.
 
-    See [[arXiv:2112.10752] High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) for more details."""
+    See [[arXiv:2112.10752] High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) for more details.
+    """
 
     def __init__(
         self,

src/refiners/training_utils/trainer.py

@@ -4,6 +4,7 @@
 from typing import Any, Callable, Generic, Literal, TypeVar, cast
 
 import torch
+import time
 from loguru import logger
 from torch import Tensor, device as Device, dtype as DType, nn
 from torch.autograd import backward
@@ -40,6 +41,24 @@
 from refiners.training_utils.gradient_clipping import GradientClipping, GradientClippingConfig
 
 
+# Ported from open-muse
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.avg: float = 0
+        self.sum: float = 0
+        self.count: int = 0
+
+    def update(self, val: float):
+        self.sum += val
+        self.count += 1
+        self.avg = self.sum / self.count
+
+
 class WarmupScheduler(LRScheduler):
     _step_count: int  # defined by LRScheduler
 
@@ -141,6 +160,10 @@ def __init__(self, config: ConfigType) -> None:
         self._models: ModelRegistry = {}
         self._callbacks: CallbackRegistry = {}
         self.config = config
+        self.batch_time_m = AverageMeter()
+        self.forward_time_m = AverageMeter()
+        self.backprop_time_m = AverageMeter()
+        self.data_time_m = AverageMeter()
         self._load_callbacks()
         self._call_callbacks(event_name="on_init_begin")
         self._load_models()
@@ -341,7 +364,7 @@ def compute_loss(self, batch: Batch) -> Tensor:
     def compute_evaluation(self) -> None:
         pass
 
-    def backward(self) -> None:
+    def backward(self, start: float) -> float:
         """Backward pass on the loss."""
         self._call_callbacks(event_name="on_backward_begin")
         scaled_loss = self.loss / self.clock.num_step_per_iteration
@@ -356,25 +379,39 @@ def backward(self) -> None:
             self._call_callbacks(event_name="on_lr_scheduler_step_begin")
             self.lr_scheduler.step()
             self._call_callbacks(event_name="on_lr_scheduler_step_end")
+        backward_time = time.time()-start
         if self.clock.is_evaluation_step:
             self.evaluate()
+        return backward_time
 
-    def step(self, batch: Batch) -> None:
+    def step(self, batch: Batch) -> tuple[float, float]:
         """Perform a single training step."""
+        start = time.time()
         self._call_callbacks(event_name="on_compute_loss_begin")
         loss = self.compute_loss(batch=batch)
         self.loss = loss
+        forward_time = time.time() - start
+        self.forward_time_m.update(forward_time)
+        start = time.time()
         self._call_callbacks(event_name="on_compute_loss_end")
-        self.backward()
+        backward_time = self.backward(start)
+        self.backprop_time_m.update(backward_time)
+        return forward_time, backward_time
 
     def epoch(self) -> None:
         """Perform a single epoch."""
+        start = time.time()
         for batch in self.dataloader:
             if self.clock.done:
                 break
             self._call_callbacks(event_name="on_batch_begin")
-            self.step(batch=batch)
+            data_time = time.time() - start
+            self.data_time_m.update(data_time)
+            forward_time, backward_time = self.step(batch=batch)
             self._call_callbacks(event_name="on_batch_end")
+            batch_time = data_time + forward_time + backward_time
+            self.batch_time_m.update(batch_time)
+            start = time.time()
 
     @staticmethod
     def get_training_seed(instance: "Trainer[BaseConfig, Any]") -> int:

src/refiners/training_utils/wandb.py

@@ -105,6 +105,24 @@ def on_train_begin(self, trainer: "TrainerWithWandb") -> None:
         self.epoch_losses = []
         self.iteration_losses = []
 
+    def on_batch_end(self, trainer: "TrainerWithWandb") -> None:
+        batch_time, forward_time, backprop_time, data_time = (
+            trainer.batch_time_m.avg,
+            trainer.forward_time_m.avg,
+            trainer.backprop_time_m.avg,
+            trainer.data_time_m.avg,
+        )
+        if trainer.clock.is_evaluation_step:
+            effective_batch_size = trainer.clock.batch_size * trainer.clock.num_step_per_iteration
+            trainer.wandb_log(
+                data={
+                    "batch_time": batch_time / effective_batch_size,
+                    "forward_time": forward_time / effective_batch_size,
+                    "backprop_time": backprop_time / effective_batch_size,
+                    "data_time": data_time / effective_batch_size,
+                }
+            )
+
     def on_compute_loss_end(self, trainer: "TrainerWithWandb") -> None:
         loss_value = trainer.loss.detach().cpu().item()
         self.epoch_losses.append(loss_value)