| """Deterministic DDP batches grouped by clean-history length.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| import random |
| from collections import defaultdict |
| from typing import Iterator |
|
|
| from torch.utils.data import Sampler |
|
|
| from .dataset import FRAMES_PER_CHUNK |
|
|
|
|
| class DistributedContextBucketBatchSampler(Sampler[list[int]]): |
| def __init__( |
| self, |
| dataset, |
| *, |
| batch_size: int, |
| rank: int, |
| world_size: int, |
| seed: int = 0, |
| drop_last: bool = True, |
| ) -> None: |
| self.dataset = dataset |
| self.batch_size = int(batch_size) |
| self.rank = int(rank) |
| self.world_size = int(world_size) |
| self.seed = int(seed) |
| self.drop_last = bool(drop_last) |
| self.epoch = 0 |
| if self.batch_size <= 0 or not 0 <= self.rank < self.world_size: |
| raise ValueError("Invalid distributed bucket sampler configuration") |
| self._buckets: dict[int, list[int]] = defaultdict(list) |
| for record_index, record in enumerate(dataset.records): |
| context_frames = int( |
| record.get( |
| "context_frames", |
| int(record["chunk_id"]) * FRAMES_PER_CHUNK, |
| ) |
| ) |
| for pair_index in range(len(dataset.PAIRS)): |
| self._buckets[context_frames].append( |
| record_index * len(dataset.PAIRS) + pair_index |
| ) |
|
|
| def set_epoch(self, epoch: int) -> None: |
| self.epoch = int(epoch) |
|
|
| def _global_batches(self) -> list[list[int]]: |
| rng = random.Random(self.seed + self.epoch) |
| global_batch_size = self.batch_size * self.world_size |
| batches: list[list[int]] = [] |
| for bucket in self._buckets.values(): |
| indices = list(bucket) |
| rng.shuffle(indices) |
| if not self.drop_last and len(indices) % global_batch_size: |
| needed = global_batch_size - len(indices) % global_batch_size |
| indices.extend((indices * math.ceil(needed / len(indices)))[:needed]) |
| usable = len(indices) - len(indices) % global_batch_size |
| batches.extend( |
| indices[start : start + global_batch_size] |
| for start in range(0, usable, global_batch_size) |
| ) |
| rng.shuffle(batches) |
| return batches |
|
|
| def __iter__(self) -> Iterator[list[int]]: |
| start = self.rank * self.batch_size |
| end = start + self.batch_size |
| for global_batch in self._global_batches(): |
| yield global_batch[start:end] |
|
|
| def __len__(self) -> int: |
| global_batch_size = self.batch_size * self.world_size |
| if self.drop_last: |
| return sum( |
| len(values) // global_batch_size for values in self._buckets.values() |
| ) |
| return sum( |
| math.ceil(len(values) / global_batch_size) |
| for values in self._buckets.values() |
| ) |
|
|
