|
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)]
|
|
|
