dataloader-rs 0.3.0

High-performance DataLoader for Rust with a PyTorch-like interface
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

#![cfg(feature = "ndarray")]

use dataloader_rs::{DataLoader, Dataset, collator::DefaultCollator, error::Result};

struct NdDs(usize);

impl Dataset for NdDs {
    type Item = ndarray::Array1<i32>;

    fn get(&self, index: usize) -> Result<Self::Item> {
        Ok(ndarray::arr1(&[
            index as i32,
            index as i32 + 1,
            index as i32 + 2,
        ]))
    }

    fn len(&self) -> usize {
        self.0
    }
}

#[test]
fn dataloader_collates_ndarray_end_to_end() {
    let mut loader = DataLoader::builder(NdDs(8))
        .batch_size(4)
        .collator(DefaultCollator)
        .num_workers(2)
        .prefetch_depth(2)
        .build();

    let batches: Vec<ndarray::Array2<i32>> = loader.iter().map(|b| b.unwrap()).collect();
    assert_eq!(batches.len(), 2);
    for b in &batches {
        assert_eq!(b.shape(), &[4, 3]);
    }
}