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

//! [`DefaultCollate`] and [`DefaultCollator`], which collates any [`DefaultCollate`] item.

use crate::{collator::Collator, error::Result};

/// Type-level collation for nested structures.
pub trait DefaultCollate: Sized + Send + 'static {
    type Batch: Send + 'static;
    fn collate_items(items: Vec<Self>) -> Result<Self::Batch>;
}

/// PyTorch-like default collator for common Rust structures.
#[derive(Clone, Copy, Debug, Default)]
pub struct DefaultCollator;

impl<Item> Collator<Item> for DefaultCollator
where
    Item: DefaultCollate,
{
    type Batch = Item::Batch;

    fn collate(&self, items: Vec<Item>) -> Result<Self::Batch> {
        Item::collate_items(items)
    }
}

#[cfg(test)]
mod tests {
    use crate::collator::{Collator, DefaultCollator};

    #[test]
    fn collates_tuple_pairs() {
        let c = DefaultCollator;
        let out = c.collate(vec![(1u32, 10i32), (2, 20), (3, 30)]).unwrap();
        assert_eq!(out, (vec![1, 2, 3], vec![10, 20, 30]));
    }
}