burn_core/data/dataloader/
batch.rs

1use super::{BatchStrategy, DataLoader, DataLoaderIterator, Progress, batcher::Batcher};
2use burn_dataset::{
3    Dataset,
4    transform::{PartialDataset, ShuffledDataset},
5};
6use burn_tensor::Device;
7use rand::SeedableRng;
8use std::ops::DerefMut;
9use std::sync::Arc;
10
11/// A data loader that can be used to iterate over a dataset in batches.
12pub struct BatchDataLoader<I, O> {
13    strategy: Box<dyn BatchStrategy<I>>,
14    dataset: Arc<dyn Dataset<I>>,
15    batcher: Arc<dyn Batcher<I, O>>,
16    device: Device,
17    rng: Option<Arc<spin::Mutex<rand::rngs::StdRng>>>,
18}
19
20impl<I, O> Clone for BatchDataLoader<I, O> {
21    fn clone(&self) -> Self {
22        Self {
23            strategy: self.strategy.clone_dyn(),
24            dataset: self.dataset.clone(),
25            batcher: self.batcher.clone(),
26            device: self.device.clone(),
27            rng: self.rng.clone(),
28        }
29    }
30}
31
32impl<I, O> BatchDataLoader<I, O> {
33    /// Creates a new batch data loader.
34    ///
35    /// # Arguments
36    ///
37    /// * `strategy` - The batch strategy.
38    /// * `dataset` - The dataset.
39    /// * `batcher` - The batcher.
40    /// * `device`  - The device to use when loading a batch.
41    /// * `rng`     - The rng determining if the dataset is shuffled each time a dataloader
42    ///   iterator is created.
43    ///
44    /// # Returns
45    ///
46    /// The batch data loader.
47    pub fn new(
48        strategy: Box<dyn BatchStrategy<I>>,
49        dataset: Arc<dyn Dataset<I>>,
50        batcher: Arc<dyn Batcher<I, O>>,
51        device: Device,
52        rng: Option<rand::rngs::StdRng>,
53    ) -> Self {
54        Self {
55            strategy,
56            dataset,
57            batcher,
58            device,
59            rng: rng.map(|rng| Arc::new(spin::Mutex::new(rng))),
60        }
61    }
62}
63
64/// A data loader iterator that can be used to iterate over a data loader.
65struct BatchDataloaderIterator<I, O> {
66    current_index: usize,
67    len: usize,
68    strategy: Box<dyn BatchStrategy<I>>,
69    dataset: Arc<dyn Dataset<I>>,
70    batcher: Arc<dyn Batcher<I, O>>,
71    device: Device,
72}
73
74impl<I, O> DataLoader<O> for BatchDataLoader<I, O>
75where
76    I: Send + Sync + Clone + 'static,
77    O: Send + 'static,
78{
79    fn iter<'a>(&'a self) -> Box<dyn DataLoaderIterator<O> + 'a> {
80        // When starting a new iteration, we first check if the dataloader was created with an rng,
81        // implying that we should shuffle the dataset beforehand, while advancing the current
82        // rng to ensure that each new iteration shuffles the dataset differently.
83        let dataset = match &self.rng {
84            Some(rng) => Arc::new(ShuffledDataset::new(
85                self.dataset.clone(),
86                rng.lock().deref_mut(),
87            )),
88            None => self.dataset.clone(),
89        };
90        Box::new(BatchDataloaderIterator::new(
91            self.strategy.clone_dyn(),
92            dataset,
93            self.batcher.clone(),
94            self.device.clone(),
95        ))
96    }
97
98    fn num_items(&self) -> usize {
99        self.dataset.len()
100    }
101
102    fn to_device(&self, device: &Device) -> Arc<dyn DataLoader<O>> {
103        let rng = self.rng.as_ref().map(|rng| {
104            let mut rng = rng.lock();
105            rng.fork()
106        });
107        Arc::new(Self::new(
108            self.strategy.clone_dyn(),
109            self.dataset.clone(),
110            self.batcher.clone(),
111            device.clone(),
112            rng,
113        ))
114    }
115
116    fn slice(&self, start: usize, end: usize) -> Arc<dyn DataLoader<O>> {
117        let rng = self.rng.as_ref().map(|rng| {
118            let mut rng = rng.lock();
119            rng.fork()
120        });
121        let dataloader = Self::new(
122            self.strategy.clone_dyn(),
123            Arc::new(PartialDataset::new(self.dataset.clone(), start, end)),
124            self.batcher.clone(),
125            self.device.clone(),
126            rng,
127        );
128        Arc::new(dataloader)
129    }
130}
131
132impl<I, O> BatchDataloaderIterator<I, O> {
133    /// Creates a new batch data loader iterator.
134    ///
135    /// # Arguments
136    ///
137    /// * `strategy` - The batch strategy.
138    /// * `dataset` - The dataset.
139    /// * `batcher` - The batcher.
140    /// * `device`  - The device to use when loading a batch.
141    ///
142    /// # Returns
143    ///
144    /// The batch data loader iterator.
145    pub fn new(
146        strategy: Box<dyn BatchStrategy<I>>,
147        dataset: Arc<dyn Dataset<I>>,
148        batcher: Arc<dyn Batcher<I, O>>,
149        device: Device,
150    ) -> Self {
151        let len = dataset.len();
152
153        BatchDataloaderIterator {
154            current_index: 0,
155            len,
156            strategy,
157            dataset,
158            batcher,
159            device,
160        }
161    }
162}
163
164impl<I, O> Iterator for BatchDataloaderIterator<I, O> {
165    type Item = Result<O, burn_dataset::DatasetError>;
166
167    fn next(&mut self) -> Option<Self::Item> {
168        while self.current_index < self.len {
169            let chunk_size = self
170                .strategy
171                .batch_size()
172                .unwrap_or(1)
173                .min(self.len - self.current_index);
174            let indexes = (self.current_index..self.current_index + chunk_size).collect();
175
176            let items = match self.dataset.get_many(indexes) {
177                Ok(items) => items,
178                Err(err) => return Some(Err(err)),
179            };
180            self.current_index += chunk_size;
181
182            for item in items {
183                self.strategy.add(item);
184            }
185
186            if let Some(items) = self.strategy.batch(false) {
187                return Some(Ok(self.batcher.batch(items, &self.device)));
188            }
189        }
190
191        if let Some(items) = self.strategy.batch(true) {
192            return Some(Ok(self.batcher.batch(items, &self.device)));
193        }
194
195        None
196    }
197}
198
199impl<I, O> DataLoaderIterator<O> for BatchDataloaderIterator<I, O> {
200    fn progress(&self) -> Progress {
201        let unit: Option<String> = Some("items".to_string());
202
203        Progress::new(self.current_index, self.len, unit)
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use std::collections::HashSet;
210
211    use super::*;
212    use crate::data::dataloader::FixBatchStrategy;
213    use crate::data::dataloader::batcher::TestBatcher;
214    use crate::data::dataset::FakeDataset;
215
216    #[test]
217    fn test_batch_dataloader() {
218        let batcher = Arc::new(TestBatcher::new());
219        let dataset = Arc::new(FakeDataset::<String>::new(27));
220        let dataloader = BatchDataLoader::new(
221            Box::new(FixBatchStrategy::new(5)),
222            dataset.clone(),
223            batcher,
224            Default::default(),
225            None,
226        );
227
228        let mut items_dataset = HashSet::new();
229        let mut items_dataloader = HashSet::new();
230
231        for item in dataset.iter().map(Result::unwrap) {
232            items_dataset.insert(item);
233        }
234
235        for items in dataloader.iter().map(Result::unwrap) {
236            for item in items {
237                items_dataloader.insert(item);
238            }
239        }
240
241        assert_eq!(items_dataset, items_dataloader);
242    }
243
244    #[test]
245    fn test_batch_dataloader_slice() {
246        let batcher = Arc::new(TestBatcher::new());
247        let dataset = Arc::new(FakeDataset::<String>::new(27));
248        let dataloader = BatchDataLoader::new(
249            Box::new(FixBatchStrategy::new(5)),
250            dataset.clone(),
251            batcher,
252            Default::default(),
253            None,
254        );
255        let dataloader_slice = dataloader.slice(5, 15);
256
257        let mut items_dataloader = HashSet::new();
258        let mut items_dataloader_slice = HashSet::new();
259
260        let mut idx = 0;
261        for items in dataloader.iter().map(Result::unwrap) {
262            for item in items {
263                if (5..15).contains(&idx) {
264                    items_dataloader.insert(item);
265                }
266                idx += 1;
267            }
268        }
269
270        for items in dataloader_slice.iter().map(Result::unwrap) {
271            for item in items {
272                items_dataloader_slice.insert(item);
273            }
274        }
275
276        assert_eq!(items_dataloader, items_dataloader_slice);
277    }
278
279    #[test]
280    fn test_batch_dataloader_incomplete_last_batch() {
281        let batcher = Arc::new(TestBatcher::new());
282        let dataset = Arc::new(FakeDataset::<String>::new(27));
283        let dataloader = BatchDataLoader::new(
284            Box::new(FixBatchStrategy::new(5)),
285            dataset,
286            batcher,
287            Default::default(),
288            None,
289        );
290
291        let batch_sizes: Vec<usize> = dataloader
292            .iter()
293            .map(Result::unwrap)
294            .map(|items| items.len())
295            .collect();
296
297        assert_eq!(batch_sizes, vec![5, 5, 5, 5, 5, 2]);
298    }
299
300    #[test]
301    fn test_batch_dataloader_exact_multiple_batches() {
302        let batcher = Arc::new(TestBatcher::new());
303        let dataset = Arc::new(FakeDataset::<String>::new(25));
304        let dataloader = BatchDataLoader::new(
305            Box::new(FixBatchStrategy::new(5)),
306            dataset,
307            batcher,
308            Default::default(),
309            None,
310        );
311
312        let batch_sizes: Vec<usize> = dataloader
313            .iter()
314            .map(Result::unwrap)
315            .map(|items| items.len())
316            .collect();
317
318        assert_eq!(batch_sizes, vec![5, 5, 5, 5, 5]);
319    }
320
321    #[test]
322    fn test_batch_dataloader_empty_dataset() {
323        let batcher = Arc::new(TestBatcher::new());
324        let dataset = Arc::new(FakeDataset::<String>::new(0));
325        let dataloader = BatchDataLoader::new(
326            Box::new(FixBatchStrategy::new(5)),
327            dataset,
328            batcher,
329            Default::default(),
330            None,
331        );
332
333        assert!(dataloader.iter().next().is_none());
334    }
335}
burn_core/data/dataloader/batch.rs

burn_core/data/dataloader/
batch.rs
