Add causal person fusion and MuVAP stage-3 baselines

third_party/MuVAP/data/packed_asd_dm.py

@@ -0,0 +1,110 @@
+"""Lightning data module for canonical packed ASD datasets."""
+
+import torch.distributed as dist
+from lightning.pytorch import LightningDataModule
+from torch.utils.data import DataLoader
+
+from data.augment import AugmentConfig
+from data.dataloaders.packed import (
+    ASDCollator,
+    FrameBudgetBatchSampler,
+    PackedASDDataset,
+)
+
+
+class PackedASDDataModule(LightningDataModule):
+    def __init__(self, cfg: dict, projection_window):
+        super().__init__()
+        self.cfg = cfg
+        self.projection_window = projection_window
+        asd = cfg["asd"]
+        self.num_workers = asd["num_workers"]
+        self.train_window = asd.get("train_window", 250)
+        self.frame_budget = asd.get("frame_budget", 2048)
+        self.min_batch = asd.get("min_batch", 4)
+        self.max_batch = asd.get("max_batch", 16)
+        self.seed = asd.get("seed", 0)
+        # Validation is forward-only, so it affords a larger budget than training.
+        self.val_frame_budget = asd.get("val_frame_budget", 4 * self.frame_budget)
+        self.val_max_batch = asd.get("val_max_batch", 64)
+        self.augment = AugmentConfig.from_cfg(asd.get("augment"))
+
+    def _replicas(self):
+        if dist.is_available() and dist.is_initialized():
+            return dist.get_world_size(), dist.get_rank()
+        return 1, 0
+
+    def setup(self, stage: str | None = None):
+        packed = self.cfg["packed_asd"]
+        if stage in ("fit", None):
+            # The context prefix is exactly the projection window's history, so
+            # every scored frame can see the past its target depends on.
+            self.train_dataset = PackedASDDataset(
+                packed["train"],
+                projection_window=self.projection_window,
+                window=self.train_window,
+                context=self.projection_window.hist_frames,
+                jitter=True,
+                augment=self.augment,
+            )
+        if stage in ("fit", "validate", "test", None):
+            # Validation runs whole tracks so every native frame gets a score.
+            self.val_dataset = PackedASDDataset(
+                packed["val"], projection_window=self.projection_window
+            )
+
+    def train_dataloader(self):
+        replicas, rank = self._replicas()
+        sampler = FrameBudgetBatchSampler(
+            self.train_dataset.chunk_lengths(),
+            frame_budget=self.frame_budget,
+            min_batch=self.min_batch,
+            max_batch=self.max_batch,
+            shuffle=True,
+            drop_last=True,
+            seed=self.seed,
+            num_replicas=replicas,
+            rank=rank,
+        )
+        return DataLoader(
+            self.train_dataset,
+            batch_sampler=sampler,
+            num_workers=self.num_workers,
+            collate_fn=ASDCollator(),
+            pin_memory=True,
+            persistent_workers=self.num_workers > 0,
+            prefetch_factor=2 if self.num_workers > 0 else None,
+        )
+
+    def val_dataloader(self):
+        """Batch whole tracks by grouping similar lengths and padding the rest.
+
+        Evaluation needs every frame of every track, so the batch cannot be
+        cropped to its shortest member the way training's is. Grouping by
+        length first keeps the padding that replaces it small, which is what
+        makes batching worth doing at all.
+
+        Lightning's automatic sampler injection is off because the train loader
+        shards itself, so this shards too. The last rank is padded by repeating
+        batches; the task de-duplicates by entity before scoring.
+        """
+        replicas, rank = self._replicas()
+        sampler = FrameBudgetBatchSampler(
+            self.val_dataset.chunk_lengths(),
+            frame_budget=self.val_frame_budget,
+            min_batch=1,
+            max_batch=self.val_max_batch,
+            shuffle=False,
+            drop_last=False,
+            num_replicas=replicas,
+            rank=rank,
+        )
+        return DataLoader(
+            self.val_dataset,
+            batch_sampler=sampler,
+            num_workers=self.num_workers,
+            collate_fn=ASDCollator(pad=True),
+            pin_memory=True,
+            persistent_workers=self.num_workers > 0,
+            prefetch_factor=2 if self.num_workers > 0 else None,
+        )