Expand description
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 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.
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. Benchmarks against PyTorch: BENCH.md.
Re-exports§
pub use collator::Collator;pub use collator::VecCollator;pub use dataset::Dataset;pub use dataset::IterableDataset;pub use error::Error;pub use error::Result;pub use loader::DataLoader;pub use loader::DataLoaderBuilder;pub use loader::DataLoaderIter;pub use loader::IndexedIter;pub use loader::SyncSource;pub use sampler::BatchSampler;pub use sampler::DistributedSampler;pub use sampler::RandomSampler;pub use sampler::Sampler;pub use sampler::SequentialSampler;pub use transform::Identity;pub use transform::Transform;