5 unstable releases
0.4.0 | Mar 16, 2023 |
---|---|
0.3.1 | Nov 13, 2022 |
0.3.0 | Nov 1, 2022 |
0.2.1 | Sep 27, 2022 |
0.2.0 | Sep 27, 2022 |
#122 in Machine learning
23 downloads per month
175KB
2K
SLoC
ai-dataloader
A rust port of pytorch
dataloader
library.
Note: This project is still heavily in development and is at an early stage.
Highlights
- Iterable or indexable (Map style)
DataLoader
. - Customizable
Sampler
,BatchSampler
andcollate_fn
. - Integration with
ndarray
andtch-rs
, CPU and GPU support. - Default collate function that will automatically collate most of your type (supporting nesting).
- Shuffling for iterable and indexable
DataLoader
.
More info in the documentation.
Examples
Examples can be found in the examples folder but here there is a simple one
use ai_dataloader::DataLoader;
let loader = DataLoader::builder(vec![(0, "hola"), (1, "hello"), (2, "hallo"), (3, "bonjour")]).batch_size(2).shuffle().build();
for (label, text) in &loader {
println!("Label {label:?}");
println!("Text {text:?}");
}
tch-rs
integration
In order to collate your data into torch tensor that can run on the GPU, you must activate the torch feature.
This feature relies on the tch crate for bindings to the C++ libTorch
API. The libtorch
library is required can be downloaded either automatically or manually. The following provides a reference on how to set up your environment to use these bindings, please refer to the tch for detailed information or support.
We advise doing the manual installation, as doctest don't pass with the automatic one.
Next Features
This features could be added in the future:
- collect function as a closure
RandomSampler
with replacement- parallel
dataloader
(using rayon?) - distributed
dataloader
Dependencies
~2.5–4.5MB
~88K SLoC