dataloader-rs 0.3.0

//! PyTorch's `DataLoader` model — dataset, sampler, collator, loader — as a typed Rust library.
//!
//! A [`DataLoader`] owns a [`Dataset`], a [`Sampler`], a [`Transform`] and a [`Collator`]: the
//! dataset only fetches items, the transform (none by default) turns each one into what the
//! collator batches. Each [`iter`](DataLoader::iter) runs one epoch and yields `Result<Batch>` in
//! epoch order, loading `num_workers` batches in parallel on the loader's thread pool, with
//! `item_concurrency` items of each batch fetched in parallel.
//!
//! ```rust
//! use dataloader_rs::{DataLoader, Dataset, RandomSampler, error::Result};
//!
//! struct Squares(usize);
//!
//! impl Dataset for Squares {
//!     type Item = u64;
//!     fn get(&self, index: usize) -> Result<u64> {
//!         Ok((index as u64).pow(2))
//!     }
//!     fn len(&self) -> usize {
//!         self.0
//!     }
//! }
//!
//! let mut loader = DataLoader::builder(Squares(1000))
//!     .batch_size(32)
//!     .sampler(RandomSampler::new(42))
//!     .num_workers(4)
//!     .prefetch_depth(8)
//!     .drop_last(true)
//!     .build();
//!
//! for batch in &mut loader {
//!     let batch: Vec<u64> = batch.unwrap();
//!     assert_eq!(batch.len(), 32);
//! }
//! ```
//!
//! | Module | Content |
//! |---|---|
//! | [`dataset`] | the [`Dataset`] contract: a source that fetches items |
//! | [`transform`] | [`Transform`]: per-item CPU work the loader applies, closures included |
//! | [`sampler`] | index order: sequential, random, distributed, batching |
//! | [`collator`] | batch assembly: `Vec`, PyTorch-like `DefaultCollator`, pinned tensors |
//! | `image` | image transforms: decode, color, crop, pad, resize, ... (feature `image`) |
//! | [`loader`] | [`DataLoader`], its builder and iterator, the worker threads |
//! | [`error`] | [`Error`] and [`Result`] |
//!
//! Features: `ndarray` and `torch-rs` add collation of `ndarray::Array` and `tch::Tensor`;
//! `async` adds `AsyncDataset` and `DataLoader::builder_async` (items fetched by futures on
//! tokio), with `FileDataset`; `http` adds `HttpDataset` and `object-store` adds
//! `ObjectStoreDataset` (in `dataset`); `image` adds the `image` module (pure-Rust decoders and
//! image transforms); `python` builds the PyO3 module behind the `dataloader-rs` Python package
//! and enables the four.
//!
//! Design, threads and failure handling:
//! [ARCHITECTURE.md](https://github.com/AzHicham/dataloader-rs/blob/dev/ARCHITECTURE.md).
//! Benchmarks against PyTorch:
//! [BENCH.md](https://github.com/AzHicham/dataloader-rs/blob/dev/BENCH.md).

// The dataset and collator are shared with worker threads through `Arc`, so
// the crate needs no unsafe code at all.
#![forbid(unsafe_code)]

pub mod collator;
pub mod dataset;
pub mod error;
#[cfg(feature = "image")]
pub mod image;
pub mod loader;
#[cfg(feature = "python")]
pub mod python;
pub mod sampler;
pub mod transform;

// ── Top-level re-exports ──────────────────────────────────────────────────────

pub use collator::{Collator, VecCollator};
#[cfg(feature = "async")]
pub use dataset::AsyncDataset;
pub use dataset::{Dataset, IterableDataset};
pub use error::{Error, Result};
#[cfg(feature = "async")]
pub use loader::{AsyncSource, MAX_RUNTIME_THREADS};
pub use loader::{DataLoader, DataLoaderBuilder, DataLoaderIter, IndexedIter, SyncSource};
pub use sampler::{BatchSampler, DistributedSampler, RandomSampler, Sampler, SequentialSampler};
pub use transform::{Identity, Transform};