diff --git a/pts/dataset/loader.py b/pts/dataset/loader.py new file mode 100644 index 0000000..f6c4fdc --- /dev/null +++ b/pts/dataset/loader.py @@ -0,0 +1,210 @@ +import itertools +from collections import defaultdict +from typing import Any, Dict, Iterable, Iterator, List, Optional # noqa: F401 + +# Third-party imports +import torch +import numpy as np + +# First-party imports +from .common import DataEntry, Dataset +from pts.feature.transform import Transformation + +DataBatch = Dict[str, Any] + + +class BatchBuffer: + def __init__( + self, batch_size: int, device: torch.device, float_type: np.dtype = np.float32 + ) -> None: + self._buffers: Dict[Any, List[Any]] = defaultdict(list) + self.batch_size = batch_size + self._size = 0 + self.device = device + self.float_type = float_type + + def add(self, d: Dict[str, List[np.ndarray]]): + if self._buffers: + assert self._buffers.keys() == d.keys() + for k, v in d.items(): + self._buffers[k].append(v) + self._size += 1 + + def __len__(self): + return self._size + + def next_batch(self) -> DataBatch: + assert self._size > 0 + n = min(self._size, self.batch_size) + batch = {k: self.stack(v[:n]) for k, v in self._buffers.items()} + for key in self._buffers.keys(): + self._buffers[key] = self._buffers[key][n:] + self._size -= n + return batch + + def stack(self, xs): + if isinstance(xs[0], np.ndarray): + data = np.asarray(xs) + if data.dtype.kind == "f": + data = data.astype(self.float_type) + return torch.from_numpy(data).to(device=self.device, non_blocking=True) + elif isinstance(xs[0], torch.Tensor): + return torch.stack(*xs) + else: + return xs # stack all other types as list + + def shuffle(self): + perm = np.random.permutation(self._size) + for key in self._buffers.keys(): + li = self._buffers[key] + self._buffers[key] = [li[i] for i in perm] + + +class DataLoader(Iterable[DataEntry]): + """ + An abstract Iterable type for iterating and transforming a dataset, + in batches of a prescribed size. + + Parameters + ---------- + dataset + The dataset from which to load data. + transform + A transformation to apply to each entry in the dataset. + batch_size + The size of the batches to emit. + device + device to use to store data on. + float_type + Floating point type to use. + """ + + def __init__( + self, + dataset: Dataset, + transform: Transformation, + batch_size: int, + device: torch.device, + float_type: np.dtype = np.float32, + ) -> None: + self.dataset = dataset + self.transform = transform + self.batch_size = batch_size + self.device = device + self.float_type = float_type + + +class TrainDataLoader(DataLoader): + """ + An Iterable type for iterating and transforming a dataset, in batches of a + prescribed size, until a given number of batches is reached. + + The transformation are applied with in training mode, i.e. with the flag + `is_train = True`. + + Parameters + ---------- + dataset + The dataset from which to load data. + transform + A transformation to apply to each entry in the dataset. + batch_size + The size of the batches to emit. + device + device to use to store data on. + num_batches_per_epoch + Number of batches to return in one complete iteration over this object. + float_type + Floating point type to use. + """ + + def __init__( + self, + dataset: Dataset, + transform: Transformation, + batch_size: int, + device: torch.device, + num_batches_per_epoch: int, + float_type: np.dtype = np.float32, + shuffle_for_training: bool = True, + num_batches_for_shuffling: int = 10, + ) -> None: + super().__init__(dataset, transform, batch_size, device, float_type) + self.num_batches_per_epoch = num_batches_per_epoch + self.shuffle_for_training = shuffle_for_training + self._num_buffered_batches = ( + num_batches_for_shuffling if shuffle_for_training else 1 + ) + self._cur_iter: Optional[Iterator] = None + self._buffer = BatchBuffer(self.batch_size, device, float_type) + + def _emit_batches_while_buffer_larger_than(self, thresh) -> Iterator[DataBatch]: + if self.shuffle_for_training: + self._buffer.shuffle() + while len(self._buffer) > thresh: + yield self._buffer.next_batch() + + def _iterate_forever(self, collection: Iterable[DataEntry]) -> Iterator[DataEntry]: + # iterate forever over the collection, the collection must be non empty + while True: + try: + first = next(iter(collection)) + except StopIteration: + raise Exception("empty dataset") + else: + for x in itertools.chain([first], collection): + yield x + + def __len__(self) -> int: + return self.num_batches_per_epoch + + def __iter__(self) -> Iterator[DataBatch]: + batch_count = 0 + if self._cur_iter is None: + self._cur_iter = self.transform( + self._iterate_forever(self.dataset), is_train=True + ) + assert self._cur_iter is not None + while True: + data_entry = next(self._cur_iter) + self._buffer.add(data_entry) + if len(self._buffer) >= self._num_buffered_batches * self.batch_size: + for batch in self._emit_batches_while_buffer_larger_than( + self.batch_size - 1 + ): + yield batch + batch_count += 1 + if batch_count >= self.num_batches_per_epoch: + return + + +class InferenceDataLoader(DataLoader): + """ + An Iterable type for iterating and transforming a dataset just once, in + batches of a prescribed size. + + The transformation are applied with in inference mode, i.e. with the flag + `is_train = False`. + + Parameters + ---------- + dataset + The dataset from which to load data. + transform + A transformation to apply to each entry in the dataset. + batch_size + The size of the batches to emit. + device + device to use to store data on. + float_type + Floating point type to use. + """ + + def __iter__(self) -> Iterator[DataBatch]: + buffer = BatchBuffer(self.batch_size, self.device, self.float_type) + for data_entry in self.transform(iter(self.dataset), is_train=False): + buffer.add(data_entry) + if len(buffer) >= self.batch_size: + yield buffer.next_batch() + if len(buffer) > 0: + yield buffer.next_batch()