burn_core/data/dataloader/
base.rs

1pub use crate::data::dataset::{Dataset, DatasetIterator};
2use core::iter::Iterator;
3
4/// A progress struct that can be used to track the progress of a data loader.
5#[derive(new, Clone, Debug)]
6pub struct Progress {
7    /// The number of items that have been processed.
8    pub items_processed: usize,
9
10    /// The total number of items that need to be processed.
11    pub items_total: usize,
12}
13
14/// A data loader iterator that can be used to iterate over a data loader.
15pub trait DataLoaderIterator<O>: Iterator<Item = O> {
16    /// Returns the progress of the data loader.
17    fn progress(&self) -> Progress;
18}
19
20/// A data loader that can be used to iterate over a dataset.
21pub trait DataLoader<O>: Send {
22    /// Returns a boxed [iterator](DataLoaderIterator) to iterate over the data loader.
23    fn iter<'a>(&'a self) -> Box<dyn DataLoaderIterator<O> + 'a>;
24    /// The number of items (not the number of batches nor the number of iterations),
25    /// corresponding to the items_total of the progress returned by the iterator.
26    fn num_items(&self) -> usize;
27}
28
29/// A super trait for [dataloader](DataLoader) that allows it to be cloned dynamically.
30///
31/// Any dataloader that implements [Clone] should also implement this automatically.
32pub trait DynDataLoader<O>: DataLoader<O> {
33    /// Clone the dataloader and returns a new one.
34    fn clone_dyn(&self) -> Box<dyn DynDataLoader<O>>;
35}
36
37impl<D, O> DynDataLoader<O> for D
38where
39    D: DataLoader<O> + Clone + 'static,
40{
41    fn clone_dyn(&self) -> Box<dyn DynDataLoader<O>> {
42        Box::new(self.clone())
43    }
44}