|
import bisect |
|
import random |
|
from typing import Iterable |
|
|
|
from torch.utils.data import Dataset, IterableDataset |
|
|
|
|
|
class ConcatRepeatDataset(Dataset): |
|
datasets: list[Dataset] |
|
cumulative_sizes: list[int] |
|
repeats: list[int] |
|
|
|
@staticmethod |
|
def cumsum(sequence, repeats): |
|
r, s = [], 0 |
|
for dataset, repeat in zip(sequence, repeats): |
|
l = len(dataset) * repeat |
|
r.append(l + s) |
|
s += l |
|
return r |
|
|
|
def __init__(self, datasets: Iterable[Dataset], repeats: list[int]): |
|
super().__init__() |
|
|
|
self.datasets = list(datasets) |
|
self.repeats = repeats |
|
|
|
assert len(self.datasets) > 0, "datasets should not be an empty iterable" |
|
assert len(self.datasets) == len( |
|
repeats |
|
), "datasets and repeats should have the same length" |
|
|
|
for d in self.datasets: |
|
assert not isinstance( |
|
d, IterableDataset |
|
), "ConcatRepeatDataset does not support IterableDataset" |
|
|
|
self.cumulative_sizes = self.cumsum(self.datasets, self.repeats) |
|
|
|
def __len__(self): |
|
return self.cumulative_sizes[-1] |
|
|
|
def __getitem__(self, idx): |
|
dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) |
|
|
|
if dataset_idx == 0: |
|
sample_idx = idx |
|
else: |
|
sample_idx = idx - self.cumulative_sizes[dataset_idx - 1] |
|
|
|
dataset = self.datasets[dataset_idx] |
|
|
|
return dataset[sample_idx % len(dataset)] |
|
|
