File size: 2,967 Bytes
d5e0d8f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """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()
)
|