burn-dataset 0.22.0-pre.2

Library with simple dataset APIs for creating ML data pipelines
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
46
47
48

use crate::dataset::{Dataset, DatasetError};
use std::error::Error;
use std::iter::Iterator;

/// Dataset iterator.
pub struct DatasetIterator<'a, I, E = DatasetError>
where
    E: Error + Send + Sync + 'static,
{
    current: usize,
    dataset: &'a dyn Dataset<I, E>,
    len: usize,
}

impl<'a, I, E> DatasetIterator<'a, I, E>
where
    E: Error + Send + Sync + 'static,
{
    /// Creates a new dataset iterator.
    pub fn new<D>(dataset: &'a D) -> Self
    where
        D: Dataset<I, E>,
    {
        DatasetIterator {
            current: 0,
            dataset,
            len: dataset.len(),
        }
    }
}

impl<I, E> Iterator for DatasetIterator<'_, I, E>
where
    E: Error + Send + Sync + 'static,
{
    type Item = Result<I, E>;

    fn next(&mut self) -> Option<Result<I, E>> {
        if self.current >= self.len {
            return None;
        }

        let index = self.current;
        self.current += 1;

        Some(self.dataset.get(index))
    }
}