dataloader-rs for Python

A DataLoader for Python with the torch.utils.data.DataLoader interface and a Rust core. Workers are threads, not processes: samples are handed over by reference instead of being pickled, an epoch starts in milliseconds, and against PyTorch it is 4.9–6.6× faster on the loader's own overhead, 3.2–3.9× on network-bound datasets and 2.5–3.5× on array samples with free-threaded Python — with no PyTorch dependency.

PyTorch's arguments: batch_size, shuffle, sampler, num_workers, collate_fn and drop_last, plus seed for reproducible shuffling, prefetch_depth, item_concurrency to fetch the items of a batch concurrently, and transform for per-item work.
Any batch type: lists by default, NumPy arrays or tensors from your collate_fn.
Batched fetching: __getitems__(indices) is used when your dataset defines it.
Async datasets: async def __getitem__ on a PyAsyncDataset, or the built-in HttpDataset, FileDataset and ObjectStoreDataset, which fetch in Rust without the GIL. Many requests are in flight at once (64, say) without a thread for each.
Image transforms in Rust: dataloader_rs.transforms (Decode, PadTo, Resize, Normalize, …, Compose) decode and transform images on the loader's threads without the GIL, with the same values as Pillow and torchvision; ImageCollator batches them in one buffer that torch.from_dlpack takes without a copy.
Exceptions keep their type: a KeyError in __getitem__ reaches you as a KeyError.
Free-threaded Python: on CPython 3.14t and 3.15t, worker threads run Python code in parallel. There are wheels for the GIL builds (3.11 to 3.15) too; there, threads speed up the work that releases the GIL (see below).

When PyTorch is faster. The workers are threads in one process, not processes. On regular CPython (the GIL build, such as 3.13), only one thread runs Python code at a time: a __getitem__ that spends its time in pure Python (tokenising text, parsing records, Python loops) gets no faster with more workers, while PyTorch's worker processes each run their own interpreter and are 2.4–2.6× faster on such datasets. Threads win when the work releases the GIL (file and network I/O, NumPy, image decoding), and on free-threaded CPython 3.14t / 3.15t. Even there, threads that all read one shared Python dict or list contend on its reference count, and PyTorch is 2.5–2.9× faster: give each thread its own copy (threading.local, see Workers, threads and the GIL).

Requires Python ≥ 3.11.

Installation
Quick start
Datasets
PyDataloader arguments
Batches and collate_fn
Transforms
Image transforms in Rust
Samplers
Workers, threads and the GIL
I/O-bound datasets
Async datasets
Errors
Coming from torch.utils.data.DataLoader
Performance

Installation

pip install dataloader-rs

Wheels	Platforms
CPython 3.11 to 3.15, and the free-threaded 3.14t and 3.15t	Linux x86_64 (a CPU with AVX2: Intel from 2013, AMD from 2017) and arm64 (glibc 2.28 or later); macOS arm64 (11.0 or later)

On other interpreters and platforms, pip builds the package from the source distribution, which needs a Rust toolchain. From a checkout of the repository, run pip install maturin && maturin develop --release.

Quick start

import numpy as np
from dataloader_rs import PyDataloader, PyDataset


class Images(PyDataset):
    def __init__(self, paths: list[str]) -> None:
        super().__init__()
        self.paths = paths

    def __len__(self) -> int:
        return len(self.paths)

    def __getitem__(self, index: int) -> np.ndarray:
        return np.load(self.paths[index])


loader = PyDataloader(
    Images(paths),
    batch_size=32,
    shuffle=True,
    seed=0,  # reproducible order, a new permutation every epoch
    num_workers=4,  # 4 worker threads
    prefetch_depth=4,  # load up to 4 batches ahead of the 4 in progress
    collate_fn=np.stack,  # one (32, ...) array per batch
    drop_last=True,
)

for epoch in range(10):
    for batch in loader:  # iterating again starts a new epoch
        ...

Datasets

A dataset is a source: it only fetches items, synchronously (PyDataset) or asynchronously (PyAsyncDataset and the Rust datasets, see Async datasets). What is done to each item afterwards (decoding, augmentation) is the loader's transform (see Transforms).

Subclass PyDataset and implement __len__ and __getitem__(index). Call super().__init__() if you define __init__.

Optionally implement __getitems__(indices) — PyTorch's batched-fetch protocol — to fetch a whole batch in one call: one database query, or one NumPy indexing, instead of batch_size Python calls. It receives a list[int] and must return one sample per index, in order, as any iterable; a wrong number of samples raises ValueError. Setting __getitems__ = None in a subclass turns it off, as in PyTorch.

class Rows(PyDataset):
    def __init__(self, table: np.ndarray) -> None:
        super().__init__()
        self.table = table

    def __len__(self) -> int:
        return len(self.table)

    def __getitem__(self, index: int) -> np.ndarray:
        return self.table[index]

    def __getitems__(self, indices: list[int]) -> list[np.ndarray]:
        return list(self.table[indices])  # one vectorised read per batch

With num_workers > 0, __getitem__ and __getitems__ run on several threads at once: keep the dataset read-only, or protect the state it changes with a lock.

PyDataloader arguments

Argument	Default	Meaning
`dataset`	—	The source of the items: a `PyDataset`, a `PyAsyncDataset`, or a Rust dataset (`HttpDataset`, `FileDataset`, `ObjectStoreDataset`; see Async datasets).
`batch_size`	`1`	The number of samples per batch (at least 1).
`shuffle`	`False`	Draws a new random order every epoch. It cannot be combined with `sampler`.
`seed`	`None`	The seed of `shuffle=True` (and only allowed with it): two loaders with the same seed produce the same sequence of epoch orders.
`sampler`	`None`	Any iterable of `int` indices, read at the start of each epoch.
`num_workers`	`0`	The number of batches loaded in parallel. With `0`, each batch is loaded in the calling thread.
`prefetch_depth`	`2`	How far ahead the workers may load: at most `num_workers + prefetch_depth` batches are in progress or waiting. It must be at least 1.
`transform`	`None`	A function `(index, sample) -> sample`, applied to every item before `collate_fn` (see Transforms).
`collate_fn`	`None`	A function `list[sample] -> batch`. Without it, a batch is a `list`.
`drop_last`	`False`	Drops the last batch of an epoch when it is smaller than `batch_size`.
`item_concurrency`	`1`	The number of items of one batch fetched in parallel, with one `__getitem__` call each (see I/O-bound datasets). Its former name, `items_in_parallel`, is still accepted with a `DeprecationWarning`.
`return_indices`	`False`	Each batch comes as `(batch, indices)`: `indices` is a `list` of the dataset indices of the batch's items, in the order the sampler drew them (see Per-item metadata).
`runtime_threads`	`None`	For async and Rust datasets only: the number of threads of the loader's tokio runtime, from 1 to 1024. `None` means one per CPU available to the process. It raises `ValueError` with a `PyDataset`.

len(loader) is the number of batches per epoch, and len(iterator) is the number of batches left in the current epoch.

Batches and `collate_fn`

Without collate_fn, a batch is the list of its samples. collate_fn receives that list and can return anything: np.stack and torch.stack work, and so does a function that returns a dict of arrays. It runs on the worker thread that loaded the batch (or the calling thread with num_workers=0), so heavy collation is parallel too.

Transforms

transform=f calls f(index, sample) for every item once the dataset has fetched it, before collate_fn: index is the item's dataset index, as the sampler drew it (shuffled or not), the same order as Rust's Transform::apply(index, item). It runs on the loader's threads, holding the interpreter only for the call: right after __getitem__ on the thread that fetched the item (the calling thread for a PyDataset with num_workers=0 and item_concurrency=1), or on the loader's pool as each item of an async or Rust dataset arrives. With item_concurrency > 1 the items of a batch are transformed in parallel. An exception it raises fails its batch and keeps its type.

A one-argument callable, such as a torchvision Compose, is wrapped:

from torchvision import transforms

compose = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)])
loader = PyDataloader(dataset, batch_size=32, transform=lambda i, x: compose(x))

The index lets the transform find what belongs to each item, without the dataset carrying it. Here, tiles of a large image are fetched in Rust by their coordinates, decoded on the loader's threads, and batched together with their coordinates:

import io

import numpy as np
from PIL import Image
from dataloader_rs import HttpDataset, PyDataloader

coords = np.array(tile_coords)  # (n, 2): x, y of each tile at level 0
tiles = HttpDataset.from_template(
    "https://tiles.example/images/image-1/tile?x={x}&y={y}&size=224",
    {"x": coords[:, 0], "y": coords[:, 1]},
    headers={"Authorization": f"Bearer {token}"},
)


def decode(png: bytes) -> np.ndarray:
    return np.asarray(Image.open(io.BytesIO(png)).convert("RGB"))


def stack(items):
    images, tile_coords = zip(*items)
    return np.stack(images), np.stack(tile_coords)


loader = PyDataloader(
    tiles,
    batch_size=64,
    num_workers=2,
    item_concurrency=32,
    transform=lambda i, png: (decode(png), coords[i]),
    collate_fn=stack,
)
for images, tile_coords in loader:  # (64, 224, 224, 3) and (64, 2)
    save(model(images), tile_coords)

To get the indices of each batch without a Python transform, pass return_indices=True (see Per-item metadata).

Image transforms in Rust

dataloader_rs.transforms has torchvision-like image transforms implemented in Rust (pure-Rust decoders, no C library). Pass a Compose of them as transform=: the loader runs it on its threads without the interpreter, so decoding scales with threads even on the GIL build, and with collate_fn=ImageCollator() the batch is assembled in Rust too.

from dataloader_rs import FileDataset, ImageCollator, PyDataloader
from dataloader_rs import transforms as T

tiles = T.Compose(
    [
        T.Decode(),  # PNG, JPEG, WebP, TIFF, BMP bytes → uint8 HWC
        T.ToRgb(),  # Image.convert("RGB")
        T.PadTo(224, 224),  # pad with 0 at the right and bottom, up to 224 x 224
        T.ToFloat(),  # ToTensor's values (/ 255)
        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        T.ToChw(),  # ToTensor's layout
    ]
)
loader = PyDataloader(
    FileDataset(paths),
    batch_size=64,
    num_workers=4,
    item_concurrency=4,
    transform=tiles,
    collate_fn=ImageCollator(buffers=8),  # one float32 buffer per batch, reused
)
for batch in loader:  # ImageBatch (64, 3, 224, 224)
    x = torch.from_dlpack(batch)  # no copy; .to("cuda") as usual

`dataloader_rs.transforms`	Pillow or torchvision equivalent	Input and output
`Decode(max_width=16384, max_height=16384, max_alloc=256 MiB)`	`Image.open(...)`; `decode_image`	bytes in; uint8 HWC out, with the image's own channels (1, 2, 3 or 4)
`ToRgb()`, `ToRgba()`, `ToGray()`	`convert("RGB" / "RGBA" / "L")`	uint8 HWC
`Crop(top, left, height, width)`	`F.crop`	HWC, either type
`CenterCrop(size)`	`CenterCrop`	HWC, either type
`Pad(padding, fill=0, padding_mode="constant")`	`Pad` (constant, edge, reflect, symmetric)	HWC, either type
`PadTo(width, height, fill=0, anchor="top_left")`	`F.pad` on the right and bottom; with `anchor="center"`, even padding on all sides	HWC, either type; never crops
`Resize(size, interpolation="bilinear", max_size=None, antialias=True)`	`Resize` (nearest, bilinear, bicubic, lanczos3)	uint8 HWC
`HFlip()`, `VFlip()`, `Rotate90(k=1)`	`F.hflip`, `F.vflip`, `numpy.rot90`	HWC, either type
`ToFloat(scale=True)`	the values of `ToTensor`, `ToDtype(float32, scale=True)`	uint8 in, float32 out
`Normalize(mean, std)`	`Normalize`	float32, HWC or CHW
`ToChw()`, `ToHwc()`	the axis order of `ToTensor`	changes the layout
`Compose(transforms)`	`Compose`	—

Same values: decoding (PNG, BMP, TIFF, lossless WebP), colour conversions, crops, padding, flips, rotations, ToFloat and Normalize give the same values as Pillow and torchvision, bit for bit. The chain Decode, ToRgb, PadTo, ToFloat, Normalize equals Pillow's decoding followed by padding at the right and bottom, ToTensor and Normalize, on RGB, RGBA, gray, palette and undersized tiles. JPEG decoding differs by at most a few levels (it uses another inverse DCT); Resize is within 1 level of Pillow for bilinear, a few for bicubic and Lanczos. torchvision v2's ToDtype(scale=True) multiplies by 1/255 and differs from ToTensor (and ToFloat) by one ulp on half the values.
Errors: a corrupt or truncated image, one above Decode's limits (checked from the header, before decoding: decompression bombs are refused), or a transform given an image it does not take (Normalize before ToFloat) raises ValueError for its batch, saying why.
Called directly, tiles(png_bytes) returns an Image, and np.asarray(image) is a read-only view of it. A Python transform can call them, with the interpreter released while they run.
Without ImageCollator, a batch is a list of Image objects, or whatever collate_fn makes of them.

ImageCollator(layout="NCHW", dtype=None, buffers=0) makes one ImageBatch per batch, in the "NCHW" or "NHWC" layout. dtype ("uint8" or "float32"), when given, checks the type of the items, and images of different shapes fail the batch. np.asarray(batch) and torch.from_dlpack(batch) are writable views of its buffer, without a copy, and keep it alive. buffers=n keeps the buffers of up to n dropped batches and writes the next ones into them: a fresh 38.5 MB float32 batch spends three quarters of its collation on page faults. It never waits for a buffer, so holding batches only costs memory.

Runnable versions of these examples, which work offline, are in examples/: all in Rust, Rust and torch mixed, torchvision transforms on bytes fetched in Rust, and tiles fetched by coordinates.

Back from numpy or torch. Image.from_array(obj, layout="hwc") makes an Image of a numpy array, a CPU torch tensor or anything np.asarray takes, so a Python transform can mix Rust and torch steps: Rust decode, a torch op on torch.from_dlpack(image) (a view), Image.from_array, then Rust Normalize and ImageCollator.

decode = T.Compose([T.Decode(), T.ToRgb(), T.PadTo(224, 224)])
normalize = T.Compose([T.ToFloat(), T.Normalize(MEAN, STD)])


def transform(i, png):
    chw = torch.from_dlpack(decode(png)).permute(2, 0, 1)  # no copy
    chw = augment(chw)  # any torch / torchvision op on uint8 CHW
    return normalize(Image.from_array(chw, layout="chw")), coords[i]

Values are uint8 or float32; other types raise TypeError and are never converted. The shape is (height, width) or, in layout, (height, width, channels) ("hwc") or (channels, height, width) ("chw"), with 1, 3 or 4 channels and Decode's limits (16384 × 16384, 256 MiB), else ValueError. A GPU tensor raises TypeError: move it with .cpu() first.
The image is HWC, so every transform and ImageCollator take it. A contiguous HWC array is copied once, into Rust's buffer; a CHW or strided one is made contiguous by numpy first (one more copy, several times faster than reading it strided). The interpreter is held during the copy.
Without numpy, objects with the buffer protocol and their own shape (a memoryview cast to (h, w, c)) are read, "hwc" only.

Per-item metadata

To know where each item of a batch belongs (the coordinates of a tile, a label), pass return_indices=True: each batch comes as (batch, indices), where indices is a list of the dataset indices of its items, in the order the sampler drew them. Look the metadata up with them; with NumPy, coords[indices]:

coords = np.array(tile_coords)  # (n, 2)
tiles_http = HttpDataset.from_template(
    "https://tiles.example/images/image-1/tile?x={x}&y={y}&size=224",
    {"x": coords[:, 0], "y": coords[:, 1]},
)
loader = PyDataloader(
    tiles_http,
    batch_size=64,
    num_workers=2,
    item_concurrency=32,
    transform=tiles,  # the Compose above
    collate_fn=ImageCollator(),
    return_indices=True,
)
for batch, indices in loader:  # ImageBatch, list of 64 ints
    save(model(torch.from_dlpack(batch)), coords[indices])

With a Rust dataset, the Rust transforms and ImageCollator, items are fetched, decoded and batched without the interpreter: only the loop itself runs Python, once per batch. It works the same with any dataset and collate_fn.

To mix Python and Rust work on each item, use a Python transform instead: it gets the item's index too, so it can return the metadata with the image, and ImageCollator then returns (batch, [metadata]):

loader = PyDataloader(
    tiles_http,
    batch_size=64,
    num_workers=2,
    item_concurrency=32,
    transform=lambda i, png: (tiles(png), coords[i]),  # Rust work, GIL released inside
    collate_fn=ImageCollator(),
)
for batch, batch_coords in loader:  # ImageBatch, list of 64 coordinate rows
    save(model(torch.from_dlpack(batch)), np.stack(batch_coords))

This form takes the interpreter once per item, to call the lambda.

Samplers

shuffle=True: a Rust permutation, new every epoch; with seed= two loaders produce the same sequence of orders.
sampler=: any finite iterable of indices, such as a list, an object whose __iter__ is a generator, or a weighted sampler. It is read whole at the start of each epoch (iter(loader)), so an endless sampler (itertools.cycle, the infinite samplers of iteration-based training) never starts one: give each epoch a finite sampler instead (Ctrl-C stops one that never ends). len(loader) uses len(sampler), so a sampler without __len__ makes len(loader) raise TypeError, as in PyTorch; iteration still works.

import random


class Weighted:
    def __init__(self, weights: list[float]) -> None:
        self.weights = weights

    def __len__(self) -> int:
        return len(self.weights)

    def __iter__(self):
        population = range(len(self.weights))
        yield from random.choices(population, weights=self.weights, k=len(self.weights))


loader = PyDataloader(dataset, batch_size=8, sampler=Weighted(weights))

Workers, threads and the GIL

num_workers > 0 loads that many batches in parallel on a thread pool of num_workers × item_concurrency threads, created once and reused by every epoch; they load batches ahead while your loop runs and the loader returns them in order.

Regular CPython: the threads share the GIL. Work that releases it — file reads, NumPy, image decoding, hashing, most C extensions — runs in parallel; pure-Python code in __getitem__ does not, and for such datasets PyTorch's worker processes are faster (2.4–2.6× in the benchmarks). Use free-threaded CPython, or keep the per-item work in code that releases the GIL.
Free-threaded CPython 3.14t / 3.15t: no GIL, so pure-Python code runs in parallel too.
```
uv python install 3.14t
uv run --python 3.14t python -X gil=0 train.py
```
Shared Python objects on free-threaded CPython: an object's reference count is cheap only for the thread that created it; other threads update it atomically. Threads that all read the same dict or list millions of times per second (a vocabulary, a lookup table) contend on it and stop scaling: a pure-Python tokeniser over one shared vocabulary runs slower on 8 threads than on 1 with 3.14t. Give each thread its own copy of hot lookup tables (threading.local), or on 3.15t key them with interned strings (sys.intern), which no longer carry this cost. Processes (PyTorch's workers) never share objects, so they do not see this.
Array samples on Linux: tune the allocator NumPy uses. With default settings, the memory of every batch goes back to the system and is faulted back in, which serialises worker threads. The settings are process-wide, so they are yours to make, at start-up:
```
# CPython 3.13 and 3.14t: NumPy allocates with glibc's malloc
GLIBC_TUNABLES=glibc.malloc.trim_threshold=67108864:glibc.malloc.mmap_threshold=33554432 python train.py
# CPython 3.15t: NumPy's buffers come from CPython's mimalloc
MIMALLOC_PURGE_DELAY=-1 python train.py
```
On 3.14t this is worth up to 2× for array samples with 4 workers, and the process keeps some free memory for reuse (80–100 MB in the benchmarks). With glibc, raise both thresholds: either one alone turns off glibc's adjustment of the other and is slower than the defaults. Batches of 16 MiB and more hit a mimalloc path no setting avoids on 3.15t: smaller batches are faster there.
Breaking out of a loop returns immediately; batches already being loaded finish in the background and are discarded.
Ctrl-C interrupts next() within 100 ms while it waits for workers, and the iterator stays usable. A batch fetched on your own thread (item_concurrency > 1 without workers) runs to the end first.
Any thread can create, iterate and drop a loader, one thread at a time per object: a second thread calling next() on an iterator that another is advancing gets RuntimeError.
Worker threads have an 8 MiB stack, a main thread's usual size, so a __getitem__ recurses as deep as it would on your thread.
multiprocessing with fork: a child process can use a loader its parent already iterated; the loader rebuilds its threads there.

I/O-bound datasets

By default a worker fetches the items of its batch one after another, so a dataset that waits on the network (an object store, an HTTP API, a database) keeps only num_workers requests in flight. With item_concurrency=k, up to k items of each batch are fetched concurrently, one __getitem__ call each, so num_workers × item_concurrency requests are in flight:

loader = PyDataloader(
    S3Images(keys),
    batch_size=32,
    num_workers=4,  # 4 batches in parallel
    item_concurrency=16,  # 16 items of each batch at once: 64 requests in flight
)

Waiting on I/O releases the GIL, so this works on regular CPython too. __getitems__ is not used when item_concurrency > 1: each item is its own call.

Each request in flight here holds a thread. For many requests, an async dataset holds a future instead.

Async datasets

Two kinds of datasets are fetched by futures on the loader's async runtime instead of threads. For both, num_workers is the number of batches in flight (at least one), item_concurrency the number of items of each batch in flight, so num_workers × item_concurrency requests run at once, from a few threads. The other arguments, the order of batches and the errors are the same as for PyDataset.

Setting	What it sizes	How to choose it
`num_workers × item_concurrency`	The requests in flight.	The latency times the items per second you need; it can be far more than the number of cores.
`runtime_threads`	The CPU used by the network: receiving data, TLS, copying bodies (on tokio).	The default (the CPUs available) is enough; lower it to leave cores to training.
The loader's thread pool: `num_workers × item_concurrency` threads, at most the CPUs available, with a `transform` or a Rust dataset; otherwise `num_workers`	The CPU used by `transform` and `collate_fn`.	It is the CPU-bound part.

PyAsyncDataset: subclass it and write async def __getitem__ (and optionally async def __getitems__(indices), used when item_concurrency=1). The loader runs the coroutines on an asyncio event loop of its own, on a background thread; dropping an iterator cancels the coroutines in flight. Objects tied to an event loop (a client session, open connections) must be created inside the coroutines, on the loader's loop, and released by an optional async def aclose(self): when the loader is dropped, it awaits aclose() once on its loop, then stops the loop. An exception from aclose(), or an aclose() still running after 5 seconds, is logged (logger dataloader_rs) and the loop stops anyway.

import aiohttp
from dataloader_rs import PyAsyncDataset, PyDataloader


class Tiles(PyAsyncDataset):
    def __init__(self, urls: list[str]) -> None:
        super().__init__()
        self.urls = urls
        self.session = None

    def __len__(self) -> int:
        return len(self.urls)

    async def __getitem__(self, index: int) -> bytes:
        if self.session is None:  # created on the loader's event loop
            self.session = aiohttp.ClientSession()
        async with self.session.get(self.urls[index]) as response:
            return await response.read()

    async def aclose(self) -> None:  # awaited on the loader's loop when it is dropped
        if self.session is not None:
            await self.session.close()


loader = PyDataloader(Tiles(urls), batch_size=64, num_workers=2, item_concurrency=32)

Rust datasets: items are bytes, fetched in Rust without holding the interpreter; the loader's transform turns them into what you need.

Dataset	Item `i`	Options
`HttpDataset(urls)`	The body of a GET request to `urls[i]`. One client keeps connections open.	`headers`; `timeout` (seconds per attempt, body included; 30 by default); `connect_timeout` (10); `retries` (3); `backoff`, `max_backoff`, `max_bytes`
`HttpDataset.from_template(template, columns, *, length)`	The body of a GET request to the URL made from `template` with the values of item `i`.	The options of `HttpDataset(urls)`
`FileDataset(paths)`	The content of `paths[i]`.	`timeout` (seconds per attempt; `None` by default); `retries` (0); `backoff`, `max_backoff`, `max_bytes`
`FileDataset.from_template(template, columns, *, length)`	The content of the path made from `template` with the values of item `i`.	The options of `FileDataset(paths)`
`ObjectStoreDataset(url, keys)`	The object `keys[i]` under the store and prefix of `url` (`s3://`, `gs://`, `az://`, `https://`).	`options` (`object_store` keys such as `aws_region`; credentials also come from the `AWS_`, `GOOGLE_` and `AZURE_*` environment variables); `timeout` (seconds per request; `None` means 30); `retries` (3); `backoff`, `max_backoff`, `max_bytes`
`ObjectStoreDataset.from_template(url, template, columns, *, length)`	The object whose key, under the prefix of `url`, is made from `template` with the values of item `i`.	The options of `ObjectStoreDataset(url, keys)`

Every from_template makes the locations in Rust, item by item, so no string is kept per item. length is required when there are no columns. The retry options are the same everywhere: up to retries more attempts after a transient failure, waiting backoff × 2ⁿ seconds with jitter, at most max_backoff seconds (0.1 and 5 by default).

import io

import numpy as np
from PIL import Image
from dataloader_rs import HttpDataset, PyDataloader

tiles = HttpDataset(
    [f"https://tiles.example/image-1/{x}_{y}.png" for x, y in coords],
    headers={"Authorization": f"Bearer {token}"},
    retries=5,  # 408, 429, 5xx and connection errors, with exponential backoff
)
loader = PyDataloader(
    tiles,
    batch_size=64,
    num_workers=2,
    item_concurrency=32,
    transform=lambda i, png: np.asarray(Image.open(io.BytesIO(png))),
    collate_fn=np.stack,
)

For many URLs that differ only by a few values (34,000 tiles of one large image, say), give a template and the values instead of a list of URLs. The template is parsed once; each URL is made in Rust when its item is fetched:

cols, rows = np.meshgrid(np.arange(200), np.arange(170))  # 34,000 (col, row) coordinates
tiles = HttpDataset.from_template(
    "https://tiles.example/images/{image}/tile?x={x}&y={y}&size=224&format=png",
    {"image": ["image 1.tif"] * cols.size, "x": cols.ravel() * 224, "y": rows.ravel() * 224},
    headers={"Authorization": f"Bearer {token}"},
)

{name} is the value of column name for the item, {index} the item index, {{ and }} literal braces. A column is a sequence or a numpy array (any integer dtype) of ints, written in decimal, or of strs, percent-encoded as a URL component ("image 1.tif" becomes image%201.tif; . and .. are refused). All columns have the same length, the number of items; with {index} only, pass length=. A placeholder without a column, a column the template does not use (probably a typo), unbalanced braces, or an invalid first URL raise ValueError; a column that is not ints or strs raises TypeError.

Files and objects take templates the same way. In a path or a key a str is written as it is, but a value cannot hold / (nor \ or NUL in a path) or be . or .., so it cannot leave the template's directory (ValueError); a key part cannot be empty either:

from dataloader_rs import FileDataset, ObjectStoreDataset

slides = ["a.svs", "a.svs", "b.svs"]
# /mnt/tiles/a.svs/0.png, /mnt/tiles/a.svs/1.png, /mnt/tiles/b.svs/2.png
files = FileDataset.from_template(
    "/mnt/tiles/{slide}/{index}.png",
    {"slide": slides},
    timeout=10.0,  # a hung network file system fails the attempt...
    retries=2,  # ...which is tried again
)
# s3://bucket/tiles/a.svs/0.png, ...: the template is relative to the prefix of the URL
objects = ObjectStoreDataset.from_template(
    "s3://bucket/tiles",
    "{slide}/{index}.png",
    {"slide": slides},
    options={"aws_region": "eu-west-3"},
)

What is retried depends on the dataset. HttpDataset: connection errors, timeouts and the statuses 408, 429, 500, 502, 503 and 504, honouring a Retry-After up to max_backoff; other statuses fail at once. FileDataset: attempts over timeout and the I/O errors TimedOut, Interrupted and WouldBlock; none by default. ObjectStoreDataset: the object_store crate retries by itself (5xx, 429, 408, connection errors, timeouts, within 3 minutes of the first attempt) with these settings; set its timeout with timeout=, not in options. Every item is capped at max_bytes (64 MiB by default), checked before and while reading. TLS uses the bundled Mozilla root certificates, not the system store.

The loader checks a dataset when it is built: async def methods on a PyDataset, plain ones on a PyAsyncDataset (aclose included), a class with both bases, or a missing __len__ raise TypeError. A plain __getitem__ that returns a coroutine (an async def helper called without await) raises TypeError on the first item.

Errors

An exception raised by __getitem__, __getitems__ (sync or async), transform or collate_fn reaches you with its own type and message, and the next next() continues with the following batch, with or without workers, as in PyTorch.
The Rust datasets raise OSError for a failed request or read (with the status for HTTP), TimeoutError when the last attempt ran out of time, ValueError for an item over max_bytes, IndexError for an index out of range. Messages show URLs without their query or credentials.
A __len__ or a sampler that raises, or a sampler that yields something other than a non-negative int, makes iter(loader) raise: an epoch never comes out silently empty or short.
Invalid arguments raise ValueError (batch_size=0, shuffle with sampler, seed without shuffle) or TypeError (a non-iterable sampler).

Coming from `torch.utils.data.DataLoader`

PyTorch	dataloader-rs
`torch.utils.data.Dataset`	subclass `PyDataset` (same `__len__` / `__getitem__` / `__getitems__`)
`batch_size`, `shuffle`, `sampler`, `num_workers`, `collate_fn`, `drop_last`	same
`generator=` for a reproducible shuffle	`seed=`
`prefetch_factor` (batches per worker)	`prefetch_depth` (batches in total, beyond the `num_workers` batches in progress)
worker processes, `persistent_workers`, `worker_init_fn`, `multiprocessing_context`	threads: there is nothing to start, copy or initialise per worker
a `Dataset` doing blocking I/O in worker processes	`PyAsyncDataset`, or `HttpDataset` / `FileDataset` / `ObjectStoreDataset`
transforms applied inside `__getitem__`	`transform=f`, applied by the loader to any dataset as `f(index, item)`
Pillow + torchvision image transforms	`dataloader_rs.transforms.Compose([...])`, in Rust, without the GIL
`pin_memory`, `timeout`, `batch_sampler`, `IterableDataset`	not supported
default collation into tensors	a `list`; pass a `collate_fn` such as `np.stack` or `torch.stack`

Performance

Against torch.utils.data.DataLoader with the same dataset code, on 4 cores of an Intel Core i9-13900H, with the allocator settings of Workers, threads and the GIL for both: 4.9–6.6× faster on the loader's own overhead, 3.2–3.9× on network-bound datasets (item_concurrency), first batches in milliseconds against 11–334 ms for torch's fresh worker processes, and 2.5–3.5× on array samples (file reads, NumPy augmentation) with free-threaded 3.14t / 3.15t, 1.0–1.5× on 3.13. Pure-Python CPU work with per-thread tables is about 1.2× faster on 3.14t and 3.15t; on shared objects or under the GIL, torch's processes are 2.4–2.9× faster. Fetching 2 048 tiles of 64 KiB over HTTP (10 ms each, batches of 64), a fresh loader per epoch: torch with 8 / 16 worker processes (requests) takes 3.0 / 1.6 s, HttpDataset with 64 requests in flight 0.38 s (PyAsyncDataset with aiohttp 0.45 s, requests threads 0.48 s on 3.14t and 1.04 s under the GIL), first batch 0.03 s against 0.8 s, on 3.13 and 3.14t alike. Every number, the method and the caveats: BENCH.md.

Source, issues and the Rust crate: github.com/AzHicham/dataloader-rs.

Navigation Menu

Repository navigation

README.md

README.md

dataloader-rs for Python

Contents

Installation

Quick start

Datasets

PyDataloader arguments

Batches and `collate_fn`

Transforms

Image transforms in Rust

Per-item metadata

Samplers

Workers, threads and the GIL

I/O-bound datasets

Async datasets

Errors

Coming from `torch.utils.data.DataLoader`

Performance

FilesExpand file tree

README.md

Latest commit

History

README.md

File metadata and controls

dataloader-rs for Python

Contents

Installation

Quick start

Datasets

PyDataloader arguments

Batches and collate_fn

Transforms

Image transforms in Rust

Per-item metadata

Samplers

Workers, threads and the GIL

I/O-bound datasets

Async datasets

Errors

Coming from torch.utils.data.DataLoader

Performance

Batches and `collate_fn`

Coming from `torch.utils.data.DataLoader`