ai-dataloader 0.6.2

Rust implementation to the PyTorch DataLoader
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
38
39
40
41
42
43
44
45

use ai_dataloader::iterable::DataLoader;

#[cfg(not(feature = "tch"))]
fn main() {
    let dataset = vec![
        (0, vec![1, 23, 4, 0]),
        (1, vec![4, 0, 0, 0]),
        (1, vec![8, 23, 12, 3]),
        (0, vec![2, 45, 4, 0]),
    ];
    // `Vec` implements `IntoIterator` and the `Item` yield are supported by the default collate function,
    // so no further work is required.
    let loader = DataLoader::builder(dataset).batch_size(2).shuffle().build();

    for (label, text) in &loader {
        dbg!(label);
        dbg!(text);
    }
}

#[cfg(feature = "tch")]
use ai_dataloader::collate::TorchCollate;

#[cfg(feature = "tch")]
fn main() {
    // Lets say we have tokenized the input.
    let dataset = vec![
        (0, vec![1, 23, 4, 0]),
        (1, vec![4, 0, 0, 0]),
        (1, vec![8, 23, 12, 3]),
        (0, vec![2, 45, 4, 0]),
    ];
    // `Vec` implements `IntoIterator` and the `Item` yield are supported by the default collate function,
    // so no further work is required.
    let loader = DataLoader::builder(dataset)
        .batch_size(2)
        .shuffle()
        .collate_fn(TorchCollate)
        .build();

    for (_label, text) in &loader {
        dbg!(_label);
        dbg!(text);
    }
}