tch_plus/vision/
mnist.rs

1//! The MNIST hand-written digit dataset.
2//!
3//! The files can be obtained from the following link:
4//! <http://yann.lecun.com/exdb/mnist/>
5use super::dataset::Dataset;
6use crate::{Kind, Tensor};
7use std::fs::File;
8use std::io::{self, BufReader, Read, Result};
9
10fn read_u32<T: Read>(reader: &mut T) -> Result<u32> {
11    let mut b = vec![0u8; 4];
12    reader.read_exact(&mut b)?;
13    let (result, _) =
14        b.iter().rev().fold((0u64, 1u64), |(s, basis), &x| (s + basis * u64::from(x), basis * 256));
15    Ok(result as u32)
16}
17
18fn check_magic_number<T: Read>(reader: &mut T, expected: u32) -> Result<()> {
19    let magic_number = read_u32(reader)?;
20    if magic_number != expected {
21        return Err(io::Error::new(
22            io::ErrorKind::Other,
23            format!("incorrect magic number {magic_number} != {expected}"),
24        ));
25    }
26    Ok(())
27}
28
29fn read_labels_(filename: &std::path::Path) -> Result<Tensor> {
30    let mut buf_reader = BufReader::new(File::open(filename)?);
31    check_magic_number(&mut buf_reader, 2049)?;
32    let samples = read_u32(&mut buf_reader)?;
33    let mut data = vec![0u8; samples as usize];
34    buf_reader.read_exact(&mut data)?;
35    Ok(Tensor::from_slice(&data).to_kind(Kind::Int64))
36}
37
38fn read_images_(filename: &std::path::Path) -> Result<Tensor> {
39    let mut buf_reader = BufReader::new(File::open(filename)?);
40    check_magic_number(&mut buf_reader, 2051)?;
41    let samples = read_u32(&mut buf_reader)?;
42    let rows = read_u32(&mut buf_reader)?;
43    let cols = read_u32(&mut buf_reader)?;
44    let data_len = samples * rows * cols;
45    let mut data = vec![0u8; data_len as usize];
46    buf_reader.read_exact(&mut data)?;
47    let tensor = Tensor::from_slice(&data)
48        .view((i64::from(samples), i64::from(rows * cols)))
49        .to_kind(Kind::Float);
50    Ok(tensor / 255.)
51}
52
53fn read_labels(filename: &std::path::Path) -> Result<Tensor> {
54    read_labels_(filename)
55        .map_err(|err| std::io::Error::new(err.kind(), format!("{filename:?} {err}")))
56}
57
58fn read_images(filename: &std::path::Path) -> Result<Tensor> {
59    read_images_(filename)
60        .map_err(|err| std::io::Error::new(err.kind(), format!("{filename:?} {err}")))
61}
62
63pub fn load_dir<T: AsRef<std::path::Path>>(dir: T) -> Result<Dataset> {
64    let dir = dir.as_ref();
65    let train_images = read_images(&dir.join("train-images-idx3-ubyte"))?;
66    let train_labels = read_labels(&dir.join("train-labels-idx1-ubyte"))?;
67    let test_images = read_images(&dir.join("t10k-images-idx3-ubyte"))?;
68    let test_labels = read_labels(&dir.join("t10k-labels-idx1-ubyte"))?;
69    Ok(Dataset { train_images, train_labels, test_images, test_labels, labels: 10 })
70}
tch_plus/vision/mnist.rs

tch_plus/vision/
mnist.rs
