torchdatasets

torchdatasets extends PyTorch's Dataset and IterableDataset with a functional, cached pipeline in the style of tf.data.

Example

import torchdatasets as td
import torchvision

# map a transform onto a dataset, then cache the results
images = Images("./data").map(torchvision.transforms.ToTensor()).cache()

# cache the first 1000 samples in RAM, the rest on disk
images = (
    dataset.map(torchvision.transforms.ToTensor())
           .cache(td.modifiers.UpToIndex(1000, td.cachers.Memory()))
           .cache(td.cachers.Pickle("./cache"))
)

Features

Use map, apply, reduce and filter directly on Dataset objects; cache in RAM, on disk, or via a custom cacher (partial caching supported); general maps such as Flatten and Select; and a WrapDataset adapter for torchvision datasets (ImageFolder, MNIST, CIFAR10). A Dataset subclass adds one super().__init__() call. Published on PyPI.