dataloader_rs/sampler/
batch_sampler.rs

1use crate::sampler::Sampler;
2
3/// Batches indices produced by an inner [`Sampler`].
4pub struct BatchSampler<S: Sampler> {
5    pub(crate) sampler: S,
6    batch_size: usize,
7    drop_last: bool,
8}
9
10impl<S: Sampler> BatchSampler<S> {
11    /// Build a batch sampler from an index sampler.
12    pub fn new(sampler: S, batch_size: usize, drop_last: bool) -> Self {
13        assert!(batch_size > 0, "batch_size must be > 0");
14        Self {
15            sampler,
16            batch_size,
17            drop_last,
18        }
19    }
20
21    pub fn batch_size(&self) -> usize {
22        self.batch_size
23    }
24
25    pub fn drop_last(&self) -> bool {
26        self.drop_last
27    }
28
29    pub fn set_batch_size(&mut self, batch_size: usize) {
30        assert!(batch_size > 0, "batch_size must be > 0");
31        self.batch_size = batch_size;
32    }
33
34    pub fn set_drop_last(&mut self, drop_last: bool) {
35        self.drop_last = drop_last;
36    }
37
38    /// Number of batches in one epoch, without advancing the sampler.
39    pub fn len(&self, dataset_len: usize) -> usize {
40        self.num_batches(self.sampler.len(dataset_len))
41    }
42
43    /// Number of batches `n_indices` sampled indices are grouped into.
44    pub fn num_batches(&self, n_indices: usize) -> usize {
45        if self.drop_last {
46            n_indices / self.batch_size
47        } else {
48            n_indices.div_ceil(self.batch_size)
49        }
50    }
51
52    /// Return one epoch of grouped batch indices.
53    pub fn batch_indices(&mut self, dataset_len: usize) -> Vec<Vec<usize>> {
54        let indices = self.sampler.indices(dataset_len);
55        let raw = indices.chunks(self.batch_size);
56        if self.drop_last {
57            raw.filter(|chunk| chunk.len() == self.batch_size)
58                .map(<[usize]>::to_vec)
59                .collect()
60        } else {
61            raw.map(<[usize]>::to_vec).collect()
62        }
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use crate::sampler::{BatchSampler, Sampler, SequentialSampler};
69
70    #[test]
71    fn batch_sampler_respects_batch_size() {
72        let mut bs = BatchSampler::new(SequentialSampler, 3, false);
73        let batches = bs.batch_indices(8);
74        assert_eq!(batches, vec![vec![0, 1, 2], vec![3, 4, 5], vec![6, 7]]);
75    }
76
77    #[test]
78    fn batch_sampler_drop_last() {
79        let mut bs = BatchSampler::new(SequentialSampler, 3, true);
80        let batches = bs.batch_indices(8);
81        assert_eq!(batches, vec![vec![0, 1, 2], vec![3, 4, 5]]);
82    }
83
84    #[test]
85    fn batch_sampler_uses_inner_sampler_order() {
86        struct ReverseSampler;
87        impl Sampler for ReverseSampler {
88            fn indices(&mut self, dataset_len: usize) -> Vec<usize> {
89                (0..dataset_len).rev().collect()
90            }
91        }
92
93        let mut bs = BatchSampler::new(ReverseSampler, 2, false);
94        let batches = bs.batch_indices(5);
95        assert_eq!(batches, vec![vec![4, 3], vec![2, 1], vec![0]]);
96    }
97
98    #[test]
99    fn batch_sampler_setters_work() {
100        let mut bs = BatchSampler::new(SequentialSampler, 2, false);
101        assert_eq!(bs.batch_indices(5), vec![vec![0, 1], vec![2, 3], vec![4]]);
102        bs.set_drop_last(true);
103        assert_eq!(bs.batch_indices(5), vec![vec![0, 1], vec![2, 3]]);
104        bs.set_batch_size(3);
105        assert_eq!(bs.batch_size(), 3);
106    }
107}
dataloader_rs/sampler/batch_sampler.rs

dataloader_rs/sampler/
batch_sampler.rs
