diff --git a/candle-nn/src/lib.rs b/candle-nn/src/lib.rs index 0eb2d8e1..2738e95f 100644 --- a/candle-nn/src/lib.rs +++ b/candle-nn/src/lib.rs @@ -8,6 +8,7 @@ pub mod layer_norm; pub mod linear; pub mod optim; pub mod var_builder; +pub mod vision; pub use activation::Activation; pub use conv::{Conv1d, Conv1dConfig}; diff --git a/candle-nn/src/vision/cifar.rs b/candle-nn/src/vision/cifar.rs new file mode 100644 index 00000000..0683c4d2 --- /dev/null +++ b/candle-nn/src/vision/cifar.rs @@ -0,0 +1,62 @@ +//! The CIFAR-10 dataset. +//! +//! The files can be downloaded from the following page: +//! +//! The binary version of the dataset is used. +use crate::vision::Dataset; +use candle::{DType, Device, Result, Tensor}; +use std::fs::File; +use std::io::{BufReader, Read}; + +const W: usize = 32; +const H: usize = 32; +const C: usize = 3; +const BYTES_PER_IMAGE: usize = W * H * C + 1; +const SAMPLES_PER_FILE: usize = 10000; + +fn read_file(filename: &std::path::Path) -> Result<(Tensor, Tensor)> { + let mut buf_reader = BufReader::new(File::open(filename)?); + let mut data = vec![0u8; SAMPLES_PER_FILE * BYTES_PER_IMAGE]; + buf_reader.read_exact(&mut data)?; + let mut images = vec![]; + let mut labels = vec![]; + for index in 0..SAMPLES_PER_FILE { + let content_offset = BYTES_PER_IMAGE * index; + labels.push(data[content_offset]); + images.push(&data[1 + content_offset..content_offset + BYTES_PER_IMAGE]); + } + let images: Vec = images + .iter() + .copied() + .flatten() + .copied() + .collect::>(); + let labels = Tensor::from_vec(labels, SAMPLES_PER_FILE, &Device::Cpu)?; + let images = Tensor::from_vec(images, (SAMPLES_PER_FILE, C, H, W), &Device::Cpu)?; + let images = (images.to_dtype(DType::F32)? / 255.)?; + Ok((images, labels)) +} + +pub fn load_dir>(dir: T) -> Result { + let dir = dir.as_ref(); + let (test_images, test_labels) = read_file(&dir.join("test_batch.bin"))?; + let train_images_and_labels = [ + "data_batch_1.bin", + "data_batch_2.bin", + "data_batch_3.bin", + "data_batch_4.bin", + "data_batch_5.bin", + ] + .iter() + .map(|x| read_file(&dir.join(x))) + .collect::>>()?; + let (train_images, train_labels): (Vec<_>, Vec<_>) = + train_images_and_labels.into_iter().unzip(); + Ok(Dataset { + train_images: Tensor::cat(&train_images, 0)?, + train_labels: Tensor::cat(&train_labels, 0)?, + test_images, + test_labels, + labels: 10, + }) +} diff --git a/candle-nn/src/vision/mnist.rs b/candle-nn/src/vision/mnist.rs new file mode 100644 index 00000000..2267f9a0 --- /dev/null +++ b/candle-nn/src/vision/mnist.rs @@ -0,0 +1,65 @@ +//! The MNIST hand-written digit dataset. +//! +//! The files can be obtained from the following link: +//! +use candle::{DType, Device, Result, Tensor}; +use std::fs::File; +use std::io::{self, BufReader, Read}; + +fn read_u32(reader: &mut T) -> Result { + 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(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 { + 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 { + 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>(dir: T) -> Result { + 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, + }) +} diff --git a/candle-nn/src/vision/mod.rs b/candle-nn/src/vision/mod.rs new file mode 100644 index 00000000..6ce743eb --- /dev/null +++ b/candle-nn/src/vision/mod.rs @@ -0,0 +1,12 @@ +use candle::Tensor; + +pub struct Dataset { + pub train_images: Tensor, + pub train_labels: Tensor, + pub test_images: Tensor, + pub test_labels: Tensor, + pub labels: usize, +} + +pub mod cifar; +pub mod mnist;