ruda_model/data/dataloader/
builder.rs

1use super::{
2    BatchDataLoader, BatchStrategy, DataLoader, FixBatchStrategy, MultiThreadDataLoader,
3    batcher::Batcher,
4};
5use ruda_dataset::Dataset;
6use ruda_tensor::api::backend::Backend;
7use rand::{SeedableRng, rngs::StdRng};
8use std::sync::Arc;
9
10/// A builder for data loaders.
11pub struct DataLoaderBuilder<B: Backend, I, O> {
12    strategy: Option<Box<dyn BatchStrategy<I>>>,
13    batcher: Arc<dyn Batcher<B, I, O>>,
14    num_threads: Option<usize>,
15    reuse_workers: bool,
16    batch_reading: bool,
17    shuffle: Option<u64>,
18    device: Option<B::Device>,
19}
20
21impl<B, I, O> DataLoaderBuilder<B, I, O>
22where
23    B: Backend,
24    I: Send + Sync + Clone + std::fmt::Debug + 'static,
25    O: Send + Clone + std::fmt::Debug + 'static,
26{
27    /// Creates a new data loader builder.
28    ///
29    /// # Arguments
30    ///
31    /// * `batcher` - The batcher.
32    ///
33    /// # Returns
34    ///
35    /// The data loader builder.
36    pub fn new<Bt>(batcher: Bt) -> Self
37    where
38        Bt: Batcher<B, I, O> + 'static,
39    {
40        Self {
41            batcher: Arc::new(batcher),
42            strategy: None,
43            num_threads: None,
44            reuse_workers: false,
45            batch_reading: false,
46            shuffle: None,
47            device: None,
48        }
49    }
50
51    /// Sets the batch size to a fix number.
52    ///
53    /// The [fix batch strategy](FixBatchStrategy) will be used.
54    ///
55    /// # Arguments
56    ///
57    /// * `batch_size` - The batch size.
58    ///
59    /// # Returns
60    ///
61    /// The data loader builder.
62    pub fn batch_size(mut self, batch_size: usize) -> Self {
63        self.strategy = Some(Box::new(FixBatchStrategy::new(batch_size)));
64        self
65    }
66
67    /// Sets the seed for shuffling.
68    ///
69    /// Each time the dataloader starts a new iteration, the dataset will be shuffled.
70    ///
71    /// # Arguments
72    ///
73    /// * `seed` - The seed.
74    ///
75    /// # Returns
76    ///
77    /// The data loader builder.
78    pub fn shuffle(mut self, seed: u64) -> Self {
79        self.shuffle = Some(seed);
80        self
81    }
82
83    /// Sets the number of workers.
84    ///
85    /// - `Some(0)` or `None`: the dataloader will run without work threads.
86    /// - `Some(n); n > 0`: the dataloader will run with `n` background threads.
87    ///
88    /// A 1-worker threaded dataloader will run loads in a background thread,
89    /// while a 0-worker threaded dataloader will run loads in the main thread.
90    ///
91    /// # Arguments
92    ///
93    /// * `num_workers` - The number of workers.
94    ///
95    /// # Returns
96    ///
97    /// The data loader builder.
98    pub fn num_workers(mut self, num_workers: usize) -> Self {
99        self.num_threads = Some(num_workers);
100        self
101    }
102
103    /// Enables reuse of idle background workers across iterations (default: false).
104    ///
105    /// Requires a positive [`Self::num_workers`] setting. Finishing or dropping an
106    /// iterator waits for its in-flight work; concurrent iterators use transient
107    /// workers if the reusable pool is busy. No effect for a single-threaded loader.
108    pub fn reuse_workers(mut self, enabled: bool) -> Self {
109        self.reuse_workers = enabled;
110        self
111    }
112
113    /// Enables batch dataset reads (default: false) for all workers.
114    ///
115    /// Uses the strategy's batch size, or one for an unknown or zero size. Reads
116    /// may fetch later samples before the strategy emits a batch. Missing items
117    /// or a mismatched result length panic without falling back to single reads.
118    pub fn batch_reading(mut self, enabled: bool) -> Self {
119        self.batch_reading = enabled;
120        self
121    }
122
123    /// Sets the data loader device.
124    ///
125    /// # Arguments
126    ///
127    /// * `device` - The device to use when loading a batch.
128    ///
129    /// # Returns
130    ///
131    /// The data loader builder.
132    pub fn set_device(mut self, device: B::Device) -> Self {
133        self.device = Some(device);
134        self
135    }
136
137    /// Builds the data loader.
138    ///
139    /// # Arguments
140    ///
141    /// * `dataset` - The dataset.
142    ///
143    /// # Returns
144    ///
145    /// The data loader.
146    pub fn build<D>(self, dataset: D) -> Arc<dyn DataLoader<B, O>>
147    where
148        D: Dataset<I> + 'static,
149    {
150        let dataset = Arc::new(dataset);
151
152        let device = self.device.unwrap_or_default();
153        let rng = self.shuffle.map(StdRng::seed_from_u64);
154        let strategy = match self.strategy {
155            Some(strategy) => strategy,
156            None => Box::new(FixBatchStrategy::new(1)),
157        };
158
159        if let Some(num_threads) = self.num_threads
160            && num_threads > 0
161        {
162            return Arc::new(
163                MultiThreadDataLoader::new(
164                    strategy,
165                    dataset,
166                    self.batcher,
167                    num_threads,
168                    device,
169                    rng,
170                )
171                .reuse_workers(self.reuse_workers)
172                .batch_reading(self.batch_reading),
173            );
174        }
175
176        Arc::new(
177            BatchDataLoader::new(strategy, dataset, self.batcher, device, rng)
178                .batch_reading(self.batch_reading),
179        )
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186    use crate::TestBackend;
187    use crate::data::dataset::FakeDataset;
188    use ruda_tensor::api::Device;
189
190    #[derive(new, Clone)]
191    struct TestBatcherDevice;
192
193    #[cfg(test)]
194    impl<I> Batcher<TestBackend, I, TestDevice> for TestBatcherDevice {
195        fn batch(&self, _items: Vec<I>, device: &TestDevice) -> TestDevice {
196            *device
197        }
198    }
199
200    type TestDevice = Device<TestBackend>;
201
202    #[test]
203    fn test_dataloader_no_workers() {
204        let default_device = TestDevice::default();
205        let dataloader = DataLoaderBuilder::new(TestBatcherDevice::new())
206            .batch_size(1)
207            .build(FakeDataset::<String>::new(9));
208
209        assert_eq!(dataloader.num_items(), 9);
210
211        for device in dataloader.iter() {
212            assert_eq!(device, default_device)
213        }
214    }
215
216    #[test]
217    fn test_dataloader_default_device() {
218        let default_device = TestDevice::default();
219        let dataloader = DataLoaderBuilder::new(TestBatcherDevice::new())
220            .batch_size(1)
221            .num_workers(1)
222            .build(FakeDataset::<String>::new(9));
223
224        assert_eq!(dataloader.num_items(), 9);
225
226        for device in dataloader.iter() {
227            assert_eq!(device, default_device)
228        }
229    }
230
231    #[test]
232    fn test_dataloader_slice_multi_device() {
233        let dataloader = DataLoaderBuilder::new(TestBatcherDevice::new())
234            .batch_size(1)
235            .num_workers(1)
236            .build(FakeDataset::<String>::new(11));
237
238        #[cfg(all(
239            test,
240            not(feature = "test-tch"),
241            not(feature = "test-wgpu"),
242            not(feature = "test-cuda")
243        ))]
244        // Only one device exists...
245        let (device1, device2) = (ruda_tensor_host::HostDevice, ruda_tensor_host::HostDevice);
246
247        #[cfg(all(test, feature = "test-tch"))]
248        let (device1, device2) = (
249            ruda_tensor_tch::LibTorchDevice::Cuda(0),
250            ruda_tensor_tch::LibTorchDevice::Cuda(1),
251        );
252
253        #[cfg(all(test, feature = "test-wgpu"))]
254        let (device1, device2) = (
255            ruda_tensor_wgpu::WgpuDevice::DiscreteGpu(0),
256            ruda_tensor_wgpu::WgpuDevice::DiscreteGpu(1),
257        );
258
259        #[cfg(all(test, feature = "test-cuda"))]
260        let (device1, device2) = (ruda_tensor_device::cuda::CudaDevice::new(0), ruda_tensor_device::cuda::CudaDevice::new(1));
261
262        assert_eq!(dataloader.num_items(), 11);
263        let dataloader_1 = dataloader.slice(0, 5).to_device(&device1);
264        let dataloader_2 = dataloader.slice(5, 11).to_device(&device2);
265
266        assert_eq!(dataloader_1.num_items(), 5);
267        assert_eq!(dataloader_2.num_items(), 6);
268
269        let (mut iterator_1, mut iterator_2) = (dataloader_1.iter(), dataloader_2.iter());
270
271        for _ in 0..5 {
272            assert_eq!(iterator_1.next(), Some(device1));
273            assert_eq!(iterator_2.next(), Some(device2));
274        }
275
276        assert_eq!(iterator_1.next(), None);
277        // For uneven split, the last dataloader (partial dataset) will have the remaining item
278        assert_eq!(iterator_2.next(), Some(device2));
279        assert_eq!(iterator_2.next(), None);
280    }
281}
ruda_model/data/dataloader/builder.rs

ruda_model/data/dataloader/
builder.rs
