Datasets

AodhML provides built-in dataset types for loading, batching, and transforming training data.

Dataset Type

import "std:dataset"

type Dataset<X, Y> = struct {
    data: [X],
    labels: [Y],
}

fn Dataset::len(self) -> u32 {
    return self.data.len()
}

fn Dataset::get(self, index: u32) -> (X, Y) {
    return (self.data[index], self.labels[index])
}

Batching

fn batch(dataset: Dataset<Tensor<f32>, Tensor<f32>>, size: u32) -> [(Tensor<f32>, Tensor<f32>)] {
    let batches = []
    let n = dataset.len()

    for i in 0..n/size {
        let start = i * size
        let end = start + size
        let x_batch = stack(dataset.data[start..end])
        let y_batch = stack(dataset.labels[start..end])
        batches.push((x_batch, y_batch))
    }

    return batches
}

DataLoader

// Planned for v0.3.0
let loader = DataLoader::new(
    dataset,
    batch_size: 32,
    shuffle: true,
    num_workers: 4,
)

for batch in loader {
    let (x, y) = batch
    let pred = model.forward(x)
    // ...
}