[WIP] Prefetching feature

Cargo.toml

@@ -33,6 +33,7 @@ rand = "0.8.5"
 tch = { version = "0.13.0", optional = true, features = ["download-libtorch"] }
 rayon = { version = "1.7.0", optional = true }
 once_cell = { version = "1.17.1", optional = true }
+crossbeam-channel = "0.5.8"
 
 
 [dev-dependencies]

examples/image.rs

@@ -127,6 +127,15 @@ fn main() {
         );
     }
 
+    let dataset = FaceLandmarksDataset::new(
+        "examples/image/dataset/face_landmarks.csv",
+        env::current_dir().unwrap().join("examples/image/dataset/"),
+    );
+    let loader = DataLoader::builder(dataset)
+        .batch_size(4)
+        .collate_fn(TorchCollate)
+        .build();
+
     loader
         .into_iter()
         .enumerate()

src/indexable/dataloader/builder.rs

@@ -3,6 +3,7 @@ use crate::{
     sampler::{BatchSampler, RandomSampler, Sampler, SequentialSampler},
     Dataset,
 };
+use std::cmp::max;
 
 #[cfg(feature = "rayon")]
 use crate::THREAD_POOL;
@@ -29,6 +30,8 @@ where
     #[cfg(feature = "rayon")]
     /// Number of threads to use.
     num_threads: usize,
+    /// Prefetch buffer size.
+    prefetch_size: usize,
 }
 
 // FIXME: kind of strange that we require DefaultCollatte even if in the end we may won't use it
@@ -56,6 +59,7 @@ where
             collate_fn: DefaultCollate,
             #[cfg(feature = "rayon")]
             num_threads,
+            prefetch_size: 0,
         }
     }
 }
@@ -72,7 +76,7 @@ where
     }
     /// Set the number of elements in a batch.
     pub fn batch_size(mut self, batch_size: usize) -> Self {
-        self.batch_sampler.batch_size = batch_size;
+        self.batch_sampler.batch_size = max(batch_size, 1);
         self
     }
 
@@ -83,6 +87,12 @@ where
         self
     }
 
+    /// Set the size of the prefetch buffer.
+    pub fn prefetch_size(mut self, prefetch_size: usize) -> Self {
+        self.prefetch_size = prefetch_size;
+        self
+    }
+
     /// Drop the lasts element if they don't feat into a batch. For instance if a dataset have 13
     /// samples and a `batch_size` of 5, the last 3 samples will be dropped.
     pub fn drop_last(mut self) -> Self {
@@ -102,6 +112,7 @@ where
             collate_fn,
             #[cfg(feature = "rayon")]
             num_threads: self.num_threads,
+            prefetch_size: self.prefetch_size,
         }
     }
 
@@ -122,6 +133,7 @@ where
             collate_fn: self.collate_fn,
             #[cfg(feature = "rayon")]
             num_threads: self.num_threads,
+            prefetch_size: self.prefetch_size,
         }
     }
     /// Create a `Dataloader` from a [`Builder`].
@@ -156,6 +168,7 @@ where
             dataset: self.dataset,
             batch_sampler: self.batch_sampler,
             collate_fn: self.collate_fn,
+            prefetch_size: self.prefetch_size,
         }
     }
 }