burn_core/data/dataloader/
base.rs

1use burn_tensor::Device;
2
3pub use crate::data::dataset::{Dataset, DatasetIterator};
4use burn_dataset::DatasetError;
5use core::iter::Iterator;
6use std::sync::Arc;
7
8/// A progress struct that can be used to track the progress of a data loader.
9#[derive(new, Clone, Debug)]
10pub struct Progress {
11    /// The number of items that have been processed.
12    pub items_processed: usize,
13
14    /// The total number of items that need to be processed.
15    pub items_total: usize,
16
17    /// The unit of measurement for the progress (e.g., "items", "batches").
18    pub unit: Option<String>,
19}
20
21impl Progress {
22    /// Returns whether all items have been processed.
23    pub fn is_completed(&self) -> bool {
24        self.items_processed == self.items_total
25    }
26}
27
28/// A data loader iterator that can be used to iterate over a data loader.
29///
30/// Yields `Err` instead of panicking when a dataset retrieval fails, so callers can decide
31/// how to handle it (e.g. abort the training run with a clear message).
32pub trait DataLoaderIterator<O>: Iterator<Item = Result<O, DatasetError>> {
33    /// Returns the progress of the data loader.
34    fn progress(&self) -> Progress;
35}
36
37/// A data loader that can be used to iterate over a dataset.
38pub trait DataLoader<O>: Send + Sync {
39    /// Returns a boxed [iterator](DataLoaderIterator) to iterate over the data loader.
40    fn iter<'a>(&'a self) -> Box<dyn DataLoaderIterator<O> + 'a>;
41
42    /// The number of items (not the number of batches nor the number of iterations),
43    /// corresponding to the items_total of the progress returned by the iterator.
44    fn num_items(&self) -> usize;
45
46    /// Move the data loader to the given device, ensuring the batches are assigned to the correct device.
47    fn to_device(&self, device: &Device) -> Arc<dyn DataLoader<O>>;
48
49    /// Returns a new data loader containing a subset of the data.
50    ///
51    /// The subset includes items from `start` (inclusive) to `end` (exclusive),
52    /// preserving the batch size and ordering of the original data loader.
53    ///
54    /// # Arguments
55    ///
56    /// * `start` - The starting index of the subset (inclusive).
57    /// * `end` - The ending index of the subset (exclusive).
58    ///
59    /// # Returns
60    ///
61    /// A boxed [`DataLoader`] instance containing only the specified range.
62    fn slice(&self, start: usize, end: usize) -> Arc<dyn DataLoader<O>>;
63}
burn_core/data/dataloader/base.rs

burn_core/data/dataloader/
base.rs
