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 and collate_fn.
• Integration with ndarray and tch-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