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, drop_last, plus seed for reproducible shuffling, prefetch_depth, and items_in_parallel to fetch the items of a batch concurrently.
• 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.
• 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. Wheels for the GIL builds (3.13 to 3.15) too: there, threads help 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.13.

Contents

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

Installation

pip install dataloader-rs

Other interpreters and platforms build from the source distribution, which needs a Rust toolchain; from a checkout: pip install maturin && maturin develop --release --features python.

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

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, one NumPy fancy-indexing, one Python call instead of batch_size. It receives a list[int] and must return one sample per index, in order (any iterable); a wrong count 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__ / __getitems__ run on several threads at once: keep the dataset read-only, or protect shared state with a lock.

PyDataloader arguments

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

Batches and collate_fn

Without collate_fn a batch is the list of its samples. collate_fn receives that list and returns anything: np.stack, torch.stack, 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.

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 — a list, a generator-producing object, 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 × items_in_parallel 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 (items_in_parallel > 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 items_in_parallel=k, up to k items of each batch are fetched concurrently, one __getitem__ call each, so num_workers × items_in_parallel requests are in flight:

loader = PyDataloader(
    S3Images(keys),
    batch_size=32,
    num_workers=4,  # 4 batches in parallel
    items_in_parallel=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 items_in_parallel > 1: each item is its own call.

Errors

• An exception raised by __getitem__, __getitems__ 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.
• 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 silently comes out 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

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 (items_in_parallel), 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. Every number, the method and the caveats: BENCH.md.

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