1 files changed, 65 insertions, 0 deletions
diff --git a/candle-datasets/src/vision/mnist.rs b/candle-datasets/src/vision/mnist.rs
new file mode 100644
index 00000000..2267f9a0
--- /dev/null
+++ b/candle-datasets/src/vision/mnist.rs
@@ -0,0 +1,65 @@
+//! The MNIST hand-written digit dataset.
+//!
+//! The files can be obtained from the following link:
+//! <http://yann.lecun.com/exdb/mnist/>
+use candle::{DType, Device, Result, Tensor};
+use std::fs::File;
+use std::io::{self, BufReader, Read};
+
+fn read_u32<T: Read>(reader: &mut T) -> Result<u32> {
+    let mut b = vec![0u8; 4];
+    reader.read_exact(&mut b)?;
+    let (result, _) = b.iter().rev().fold((0u64, 1u64), |(s, basis), &x| {
+        (s + basis * u64::from(x), basis * 256)
+    });
+    Ok(result as u32)
+}
+
+fn check_magic_number<T: Read>(reader: &mut T, expected: u32) -> Result<()> {
+    let magic_number = read_u32(reader)?;
+    if magic_number != expected {
+        Err(io::Error::new(
+            io::ErrorKind::Other,
+            format!("incorrect magic number {magic_number} != {expected}"),
+        ))?;
+    }
+    Ok(())
+}
+
+fn read_labels(filename: &std::path::Path) -> Result<Tensor> {
+    let mut buf_reader = BufReader::new(File::open(filename)?);
+    check_magic_number(&mut buf_reader, 2049)?;
+    let samples = read_u32(&mut buf_reader)?;
+    let mut data = vec![0u8; samples as usize];
+    buf_reader.read_exact(&mut data)?;
+    let samples = data.len();
+    Tensor::from_vec(data, samples, &Device::Cpu)
+}
+
+fn read_images(filename: &std::path::Path) -> Result<Tensor> {
+    let mut buf_reader = BufReader::new(File::open(filename)?);
+    check_magic_number(&mut buf_reader, 2051)?;
+    let samples = read_u32(&mut buf_reader)? as usize;
+    let rows = read_u32(&mut buf_reader)? as usize;
+    let cols = read_u32(&mut buf_reader)? as usize;
+    let data_len = samples * rows * cols;
+    let mut data = vec![0u8; data_len];
+    buf_reader.read_exact(&mut data)?;
+    let tensor = Tensor::from_vec(data, (samples, rows * cols), &Device::Cpu)?;
+    tensor.to_dtype(DType::F32)? / 255.
+}
+
+pub fn load_dir<T: AsRef<std::path::Path>>(dir: T) -> Result<crate::vision::Dataset> {
+    let dir = dir.as_ref();
+    let train_images = read_images(&dir.join("train-images-idx3-ubyte"))?;
+    let train_labels = read_labels(&dir.join("train-labels-idx1-ubyte"))?;
+    let test_images = read_images(&dir.join("t10k-images-idx3-ubyte"))?;
+    let test_labels = read_labels(&dir.join("t10k-labels-idx1-ubyte"))?;
+    Ok(crate::vision::Dataset {
+        train_images,
+        train_labels,
+        test_images,
+        test_labels,
+        labels: 10,
+    })
+}