mirror of
https://github.com/wassname/pytorch-ts.git
synced 2026-07-25 13:30:12 +08:00
37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
from typing import Callable, Iterable, Iterator, List, Optional
|
|
import itertools
|
|
|
|
from torch.utils.data import IterableDataset
|
|
|
|
from gluonts.dataset.common import Dataset
|
|
from gluonts.transform import Transformation, TransformedDataset
|
|
from gluonts.itertools import Cyclic, PseudoShuffled
|
|
|
|
|
|
class TransformedIterableDataset(IterableDataset):
|
|
def __init__(
|
|
self,
|
|
dataset: Dataset,
|
|
transform: Transformation,
|
|
is_train: bool = True,
|
|
shuffle_buffer_length: Optional[int] = None,
|
|
):
|
|
super().__init__()
|
|
self.shuffle_buffer_length = shuffle_buffer_length
|
|
|
|
self.transformed_dataset = TransformedDataset(
|
|
Cyclic(dataset),
|
|
transform,
|
|
is_train=is_train,
|
|
)
|
|
|
|
def __iter__(self):
|
|
if self.shuffle_buffer_length is None:
|
|
return iter(self.transformed_dataset)
|
|
else:
|
|
shuffled = PseudoShuffled(
|
|
self.transformed_dataset,
|
|
shuffle_buffer_length=self.shuffle_buffer_length,
|
|
)
|
|
return iter(shuffled)
|