Add DeMemWM rollout progress metrics

algorithms/dememwm/df_video.py

@@ -1,6 +1,7 @@
 from collections.abc import Mapping
 import os
 import random
+import time
 import numpy as np
 import torch
 import torch.distributed as dist
@@ -39,6 +40,25 @@ def _cfg_get(cfg, key: str, default=None):
     return getattr(cfg, key, default)
 
 
+def _cuda_device_index(device):
+    device = torch.device(device)
+    if device.type != "cuda" or not torch.cuda.is_available():
+        return None
+    return torch.cuda.current_device() if device.index is None else device.index
+
+
+def _cuda_vram_postfix(device):
+    device_index = _cuda_device_index(device)
+    if device_index is None:
+        return None
+
+    gib = 1024**3
+    allocated = torch.cuda.memory_allocated(device_index) / gib
+    reserved = torch.cuda.memory_reserved(device_index) / gib
+    peak = torch.cuda.max_memory_allocated(device_index) / gib
+    return f"{allocated:.1f}/{reserved:.1f}G peak {peak:.1f}G"
+
+
 def _trainability_cfg(cfg):
     return _cfg_get(cfg, "trainability", {})
 
@@ -714,8 +734,20 @@ class DeMemWMMinecraft(DiffusionForcingBase):
             xs_pred = xs.clone()
             n_context_frames = min(self.context_frames // self.frame_stack, target_length)
             curr_frame = n_context_frames
+            generated_frames = 0
+            rollout_start_time = time.perf_counter()
+            pbar = None
+            if curr_frame < target_length:
+                pbar = tqdm(
+                    total=target_length,
+                    initial=curr_frame,
+                    desc=f"{namespace} sampling[{batch_idx}]",
+                    unit="frame",
+                    dynamic_ncols=True,
+                )
 
             while curr_frame < target_length:
+                chunk_start_time = time.perf_counter()
                 horizon = min(target_length - curr_frame, self.chunk_size) if self.chunk_size > 0 else target_length - curr_frame
                 target_slice = slice(curr_frame, curr_frame + horizon)
                 xs_pred[target_slice] = torch.randn_like(xs_pred[target_slice]).clamp(-self.clip_noise, self.clip_noise)
@@ -822,7 +854,27 @@ class DeMemWMMinecraft(DiffusionForcingBase):
                     packed_latents[:target_length_step] = sampled_target[:target_length_step]
 
                 xs_pred[local_slice] = packed_latents[:target_length_step]
-                curr_frame += horizon
+                chunk_seconds = max(time.perf_counter() - chunk_start_time, 1e-9)
+                end_frame = curr_frame + horizon
+                curr_frame = end_frame
+                generated_frames += horizon
+
+                if pbar is not None:
+                    rollout_seconds = max(time.perf_counter() - rollout_start_time, 1e-9)
+                    postfix = {
+                        "range": f"{start_frame}:{end_frame}",
+                        "sec/frame": f"{chunk_seconds / horizon:.3f}",
+                        "frames/s": f"{horizon / chunk_seconds:.2f}",
+                        "avg_frames/s": f"{generated_frames / rollout_seconds:.2f}",
+                    }
+                    vram = _cuda_vram_postfix(xs.device)
+                    if vram is not None:
+                        postfix["vram"] = vram
+                    pbar.update(horizon)
+                    pbar.set_postfix(postfix)
+
+            if pbar is not None:
+                pbar.close()
 
             eval_start = n_context_frames
             latent_loss = F.mse_loss(xs_pred[eval_start:target_length], xs[eval_start:target_length], reduction="none")