burn_core/data/dataloader/
split.rs

1use std::sync::Arc;
2
3use burn_tensor::Device;
4
5use super::DataLoader;
6
7/// Splits a dataloader into multiple partial dataloaders (one per device).
8pub fn split_dataloader<O>(
9    dataloader: Arc<dyn DataLoader<O>>,
10    devices: &[Device],
11) -> Vec<Arc<dyn DataLoader<O>>> {
12    let num_splits = devices.len();
13    if num_splits > 1 {
14        let num_items = dataloader.num_items();
15        let mut dataloaders = Vec::with_capacity(num_splits);
16
17        let mut start = 0;
18        let step = num_items / num_splits;
19        for (i, device) in devices.iter().enumerate() {
20            let end = if i == (num_splits - 1) {
21                num_items
22            } else {
23                start + step
24            };
25            let dataloader = dataloader.slice(start, end).to_device(device);
26            dataloaders.push(dataloader);
27            start = end;
28        }
29        dataloaders
30    } else {
31        vec![dataloader]
32    }
33}
34
35#[cfg(test)]
36mod tests {
37    use burn_tensor::Device;
38    use std::collections::HashSet;
39
40    use super::*;
41    use crate::data::dataloader::batcher::Batcher;
42    use crate::data::dataloader::{BatchDataLoader, FixBatchStrategy};
43    use crate::data::dataset::FakeDataset;
44
45    #[test]
46    fn test_split_batch_dataloader() {
47        #[derive(new, Clone)]
48        pub struct TestBatcher;
49
50        #[cfg(test)]
51        impl<I> Batcher<I, (Vec<I>, Device)> for TestBatcher {
52            fn batch(&self, items: Vec<I>, device: &Device) -> (Vec<I>, Device) {
53                (items, device.clone())
54            }
55        }
56
57        let batcher = Arc::new(TestBatcher::new());
58        let dataset = Arc::new(FakeDataset::<String>::new(11));
59
60        #[allow(clippy::arc_with_non_send_sync)]
61        let dataloader = Arc::new(BatchDataLoader::new(
62            Box::new(FixBatchStrategy::new(5)),
63            dataset.clone(),
64            batcher,
65            Default::default(),
66            None,
67        ));
68
69        #[cfg(all(test, not(feature = "tch"), not(feature = "cuda")))]
70        // Only one device exists...
71        let (device1, device2) = (Device::flex(), Device::flex());
72
73        #[cfg(all(test, feature = "tch"))]
74        let (device1, device2) = (Device::libtorch_cuda(0), Device::libtorch_cuda(1));
75
76        #[cfg(all(test, feature = "cuda"))]
77        let (device1, device2) = (Device::cuda(0), Device::cuda(1));
78
79        let dataloaders = split_dataloader(dataloader.clone(), &[device1.clone(), device2.clone()]);
80
81        assert_eq!(dataloaders.len(), 2);
82
83        let [dataloader_1, dataloader_2] = match dataloaders.try_into() {
84            Ok(arr) => arr,
85            Err(_) => unreachable!(),
86        };
87        assert_eq!(dataloader_1.num_items(), 5);
88        assert_eq!(dataloader_2.num_items(), 6);
89
90        let mut items_dataloader = HashSet::new();
91        let mut items_dataloader_split = HashSet::new();
92
93        for (items, _device) in dataloader.iter().map(Result::unwrap) {
94            for item in items {
95                items_dataloader.insert(item);
96            }
97        }
98
99        for (items, device) in dataloader_1.iter().map(Result::unwrap) {
100            assert_eq!(&device, &device1);
101            for item in items {
102                items_dataloader_split.insert(item);
103            }
104        }
105
106        for (items, device) in dataloader_2.iter().map(Result::unwrap) {
107            assert_eq!(&device, &device2);
108            for item in items {
109                items_dataloader_split.insert(item);
110            }
111        }
112
113        assert_eq!(items_dataloader, items_dataloader_split);
114    }
115}
burn_core/data/dataloader/split.rs

burn_core/data/dataloader/
split.rs
