lumen_dataset/base/
dataloader.rs

1use rand::seq::SliceRandom;
2use super::Dataset;
3use super::Batcher;
4
5pub struct DataLoader<D, B> 
6where 
7    D: Dataset,
8    B: Batcher<Item = D::Item>
9{
10    dataset: D,
11    batcher: B,
12    batch_size: usize,
13    shuffle: bool,
14} 
15
16#[derive(Debug, thiserror::Error)]
17pub enum DataLoaderError {
18    #[error("dataset error: {0}")]
19    Dataset(String),
20
21    #[error("batcher error: {0}")]
22    Batcher(String),
23}
24
25pub struct DataLoaderIter<'a, D, B> 
26where 
27    D: Dataset,
28    B: Batcher<Item = D::Item>
29{
30    loader: &'a DataLoader<D, B>,
31    cursor: usize,
32    indices: Vec<usize>,
33}
34
35pub struct DataLoaderIntoIter<D, B> 
36where 
37    D: Dataset,
38    B: Batcher<Item = D::Item>
39{
40    loader: DataLoader<D, B>,
41    cursor: usize,
42    indices: Vec<usize>,
43}
44
45impl<D, B> DataLoader<D, B> 
46where 
47    D: Dataset,
48    B: Batcher<Item = D::Item> 
49{
50    pub fn new(dataset: D, batcher: B, batch_size: usize, shuffle: bool) -> Self {
51        Self {
52            dataset,
53            batcher,
54            batch_size,
55            shuffle,
56        }
57    }
58
59    pub fn dataset_len(&self) -> usize {
60        self.dataset.len()
61    }
62
63    pub fn batch_count(&self) -> usize {
64        self.dataset.len() / self.batch_size
65    }
66
67    pub fn batch_size(&self) -> usize {
68        self.batch_size
69    }
70
71    pub fn iter<'a>(&'a self) -> DataLoaderIter<'a, D, B> {
72        let indices = self.get_iter_indices();
73        DataLoaderIter { loader: self, cursor: 0, indices }
74    }
75
76    fn get_iter_indices(&self) -> Vec<usize> {
77        let len = self.dataset.len();
78        let mut indices: Vec<usize> = (0..len).collect();
79        
80        if self.shuffle {
81            let mut rng = rand::rng();
82            indices.shuffle(&mut rng);
83        }
84
85        indices
86    }
87}
88
89impl<D, B> DataLoader<D, B> 
90where 
91    D: Dataset,
92    B: Batcher<Item = D::Item> + Default
93{
94    pub fn from_dataset(dataset: D, batch_size: usize, shuffle: bool) -> Self {
95        Self::new(dataset, B::default(), batch_size, shuffle)
96    }
97}
98
99impl<D, B> Iterator for DataLoaderIntoIter<D, B> 
100where 
101    D: Dataset,
102    B: Batcher<Item = D::Item> 
103{
104    type Item = Result<B::Output, DataLoaderError>;
105    fn next(&mut self) -> Option<Self::Item> {
106        iter_next(&self.loader, &mut self.cursor, &self.indices)
107    }
108}
109
110impl<'a, D, B> Iterator for DataLoaderIter<'a, D, B> 
111where 
112    D: Dataset,
113    B: Batcher<Item = D::Item> 
114{
115    type Item = Result<B::Output, DataLoaderError>;
116    fn next(&mut self) -> Option<Self::Item> {
117        iter_next(&self.loader, &mut self.cursor, &self.indices)
118    }
119}
120
121impl<D, B> IntoIterator for DataLoader<D, B> 
122where 
123    D: Dataset,
124    B: Batcher<Item = D::Item> 
125{
126    type IntoIter = DataLoaderIntoIter<D, B>;
127    type Item = Result<B::Output, DataLoaderError>;
128    fn into_iter(self) -> Self::IntoIter {
129        let indices = self.get_iter_indices();
130        DataLoaderIntoIter {
131            loader: self, 
132            cursor: 0,
133            indices,
134        } 
135    }
136}
137
138fn iter_next<D, B>(loader: &DataLoader<D, B>, cursor: &mut usize, indices: &[usize]) -> Option<Result<B::Output, DataLoaderError>> 
139where 
140    D: Dataset,
141    B: Batcher<Item = D::Item> 
142{
143    let begin = *cursor;
144    if begin >= loader.dataset.len() {
145        None
146    } else {
147        let mut f = || -> Result<B::Output, DataLoaderError> {
148            let end = *cursor + loader.batch_size;
149            let end = end.min(loader.dataset.len());
150            let mut items = vec![];
151            for index in begin..end {
152                let index = indices[index];
153                let item = loader.dataset
154                    .get(index)
155                    .map_err(|e| DataLoaderError::Dataset(e.to_string()))?
156                    .expect("must valid item");
157                items.push(item);
158            }        
159            *cursor += items.len();
160            let batch = loader.batcher
161                .batch(items)
162                .map_err(|e| DataLoaderError::Batcher(e.to_string()))?;
163
164            Ok(batch)
165        };
166
167        Some(f())
168    }
169}
lumen_dataset/base/dataloader.rs

lumen_dataset/base/
dataloader.rs
