dataloader_rs/
lib.rs

1//! PyTorch's `DataLoader` model — dataset, sampler, collator, loader — as a typed Rust library.
2//!
3//! A [`DataLoader`] owns a [`Dataset`], a [`Sampler`], a [`Transform`] and a [`Collator`]: the
4//! dataset only fetches items, the transform (none by default) turns each one into what the
5//! collator batches. Each [`iter`](DataLoader::iter) runs one epoch and yields `Result<Batch>` in
6//! epoch order, loading `num_workers` batches in parallel on the loader's thread pool, with
7//! `item_concurrency` items of each batch fetched in parallel.
8//!
9//! ```rust
10//! use dataloader_rs::{DataLoader, Dataset, RandomSampler, error::Result};
11//!
12//! struct Squares(usize);
13//!
14//! impl Dataset for Squares {
15//!     type Item = u64;
16//!     fn get(&self, index: usize) -> Result<u64> {
17//!         Ok((index as u64).pow(2))
18//!     }
19//!     fn len(&self) -> usize {
20//!         self.0
21//!     }
22//! }
23//!
24//! let mut loader = DataLoader::builder(Squares(1000))
25//!     .batch_size(32)
26//!     .sampler(RandomSampler::new(42))
27//!     .num_workers(4)
28//!     .prefetch_depth(8)
29//!     .drop_last(true)
30//!     .build();
31//!
32//! for batch in &mut loader {
33//!     let batch: Vec<u64> = batch.unwrap();
34//!     assert_eq!(batch.len(), 32);
35//! }
36//! ```
37//!
38//! | Module | Content |
39//! |---|---|
40//! | [`dataset`] | the [`Dataset`] contract: a source that fetches items |
41//! | [`transform`] | [`Transform`]: per-item CPU work the loader applies, closures included |
42//! | [`sampler`] | index order: sequential, random, distributed, batching |
43//! | [`collator`] | batch assembly: `Vec`, PyTorch-like `DefaultCollator`, pinned tensors |
44//! | `image` | image transforms: decode, color, crop, pad, resize, ... (feature `image`) |
45//! | [`loader`] | [`DataLoader`], its builder and iterator, the worker threads |
46//! | [`error`] | [`Error`] and [`Result`] |
47//!
48//! Features: `ndarray` and `torch-rs` add collation of `ndarray::Array` and `tch::Tensor`;
49//! `async` adds `AsyncDataset` and `DataLoader::builder_async` (items fetched by futures on
50//! tokio), with `FileDataset`; `http` adds `HttpDataset` and `object-store` adds
51//! `ObjectStoreDataset` (in `dataset`); `image` adds the `image` module (pure-Rust decoders and
52//! image transforms); `python` builds the PyO3 module behind the `dataloader-rs` Python package
53//! and enables the four.
54//!
55//! Design, threads and failure handling:
56//! [ARCHITECTURE.md](https://github.com/AzHicham/dataloader-rs/blob/dev/ARCHITECTURE.md).
57//! Benchmarks against PyTorch:
58//! [BENCH.md](https://github.com/AzHicham/dataloader-rs/blob/dev/BENCH.md).
59
60// The dataset and collator are shared with worker threads through `Arc`, so
61// the crate needs no unsafe code at all.
62#![forbid(unsafe_code)]
63
64pub mod collator;
65pub mod dataset;
66pub mod error;
67#[cfg(feature = "image")]
68pub mod image;
69pub mod loader;
70#[cfg(feature = "python")]
71pub mod python;
72pub mod sampler;
73pub mod transform;
74
75// ── Top-level re-exports ──────────────────────────────────────────────────────
76
77pub use collator::{Collator, VecCollator};
78#[cfg(feature = "async")]
79pub use dataset::AsyncDataset;
80pub use dataset::{Dataset, IterableDataset};
81pub use error::{Error, Result};
82#[cfg(feature = "async")]
83pub use loader::{AsyncSource, MAX_RUNTIME_THREADS};
84pub use loader::{DataLoader, DataLoaderBuilder, DataLoaderIter, IndexedIter, SyncSource};
85pub use sampler::{BatchSampler, DistributedSampler, RandomSampler, Sampler, SequentialSampler};
86pub use transform::{Identity, Transform};
dataloader_rs/lib.rs

dataloader_rs/
lib.rs
