111 lines
4.2 KiB
Python
111 lines
4.2 KiB
Python
"""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,
|
|
)
|