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}