diff --git a/Cargo.toml b/Cargo.toml index 8b861e7..d9b2515 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,9 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[features] +visualization = ["dep:image", "dep:viuer"] + [dependencies] boxed-array = "0.1.0" nalgebra = { version = "^0.32", features = ["std", "macros", "rand", "serde-serialize"] } @@ -14,6 +17,8 @@ num = "^0.4" rand = "^0.8" rand_distr = "0.4.3" textplots = "0.8.0" +image = { version = "0.24.6", optional = true } +viuer = { version = "0.6.2", optional = true } [dev-dependencies] image = "0.24.6" @@ -24,3 +29,15 @@ approx = "0.5.1" [profile.release] debug = true + +[[example]] +name = "densenet" +required-features = ["visualization"] + +[[example]] +name = "mnist-decoder" +required-features = ["visualization"] + +[[example]] +name = "mnist-diffusion" +required-features = ["visualization"] diff --git a/examples/densenet.rs b/examples/densenet.rs index 9e41e46..3bcfdd8 100644 --- a/examples/densenet.rs +++ b/examples/densenet.rs @@ -51,16 +51,16 @@ fn main() { &test_inputs, ); - draw_neuron_activation( + neuramethyst::draw_neuron_activation( |input| { - let output = network.eval(&dvector![input[0] as f32, input[1] as f32]); + let output = network.eval(&dvector![input[0], input[1]]); let estimation = output[0] / (output[0] + output[1]); - let color = network.eval(&dvector![input[0] as f32, input[1] as f32]); + let color = network.eval(&dvector![input[0], input[1]]); (&color / color.map(|x| x * x).sum() * estimation) .into_iter() - .map(|x| x.abs() as f64) + .map(|x| x.abs() as f32) .collect::>() }, 6.0, @@ -96,42 +96,6 @@ fn main() { } } -// TODO: move this to the library? -fn draw_neuron_activation Vec>(callback: F, scale: f64) { - use viuer::Config; - - const WIDTH: u32 = 64; - const HEIGHT: u32 = 64; - - let mut image = image::RgbImage::new(WIDTH, HEIGHT); - - fn sigmoid(x: f64) -> f64 { - 1.9 / (1.0 + (-x * 3.0).exp()) - 0.9 - } - - for y in 0..HEIGHT { - let y2 = 2.0 * y as f64 / HEIGHT as f64 - 1.0; - for x in 0..WIDTH { - let x2 = 2.0 * x as f64 / WIDTH as f64 - 1.0; - let activation = callback([x2 * scale, y2 * scale]); - let r = (sigmoid(activation.get(0).copied().unwrap_or(-1.0)) * 255.0).floor() as u8; - let g = (sigmoid(activation.get(1).copied().unwrap_or(-1.0)) * 255.0).floor() as u8; - let b = (sigmoid(activation.get(2).copied().unwrap_or(-1.0)) * 255.0).floor() as u8; - - *image.get_pixel_mut(x, y) = image::Rgb([r, g, b]); - } - } - - let config = Config { - use_kitty: false, - truecolor: true, - // absolute_offset: false, - ..Default::default() - }; - - viuer::print(&image::DynamicImage::ImageRgb8(image), &config).unwrap(); -} - fn one_hot(value: usize, categories: usize) -> DVector { let mut res = DVector::from_element(categories, 0.0); if value < categories { diff --git a/examples/mnist-decoder.rs b/examples/mnist-decoder.rs new file mode 100644 index 0000000..2d59b5f --- /dev/null +++ b/examples/mnist-decoder.rs @@ -0,0 +1,144 @@ +use nalgebra::DVector; +use rust_mnist::Mnist; + +use neuramethyst::{ + argmax, cycle_shuffling, + derivable::{ + activation::{Linear, Logistic, Relu, Swish, Tanh}, + loss::{CrossEntropy, Euclidean}, + }, + plot_losses, + prelude::*, +}; + +const TRAIN_SIZE: usize = 50000; +const TEST_SIZE: usize = 1000; +const WIDTH: usize = 28; +const HEIGHT: usize = 28; +const LATENT_SIZE: usize = 25; + +pub fn main() { + let Mnist { + train_data: train_images, + train_labels, + test_data: test_images, + test_labels, + .. + } = Mnist::new("data/"); + + let train_images = train_images + .into_iter() + .map(|raw| { + DVector::from_iterator(WIDTH * HEIGHT, raw.into_iter().map(|x| x as f32 / 255.0)) + }) + .take(TRAIN_SIZE); + let train_labels = train_labels + .into_iter() + .map(|x| one_hot(x as usize, 10)) + .take(TRAIN_SIZE); + + let test_images = test_images + .into_iter() + .map(|raw| { + DVector::from_iterator(WIDTH * HEIGHT, raw.into_iter().map(|x| x as f32 / 255.0)) + }) + .take(TEST_SIZE); + let test_labels = test_labels + .into_iter() + .map(|x| one_hot(x as usize, 10)) + .take(TEST_SIZE); + + let test_data = test_images + .clone() + .zip(test_images.clone()) + .collect::>(); + + // First, train an encoder-decoder network (unsupervised) + + let mut network = neura_sequential![ + neura_layer!("dense", 100).activation(Swish(Logistic)), + neura_layer!("dense", 50).activation(Swish(Logistic)), + neura_layer!("dense", LATENT_SIZE).activation(Tanh), + neura_layer!("dense", 50), + neura_layer!("dense", 100), + neura_layer!("dense", WIDTH * HEIGHT).activation(Relu), + ] + .construct(NeuraShape::Vector(WIDTH * HEIGHT)) + .unwrap(); + + let trainer = NeuraBatchedTrainer::with_epochs(0.03, 75, 512, TRAIN_SIZE); + // trainer.log_iterations = 1; + + let losses = trainer.train( + &NeuraBackprop::new(Euclidean), + &mut network, + cycle_shuffling( + train_images.clone().zip(train_images.clone()), + rand::thread_rng(), + ), + &test_data, + ); + + plot_losses(losses, 128, 48); + + // Then, train a small network to decode the encoded data into the categories + + let trimmed_network = network.clone().trim_tail().trim_tail().trim_tail(); + + let mut network = neura_sequential![ + ..trimmed_network.lock(), + neura_layer!("dense", LATENT_SIZE) + .activation(Tanh) + .construct(NeuraShape::Vector(LATENT_SIZE)) + .unwrap(), + neura_layer!("dense", 10) + .activation(Linear) + .construct(NeuraShape::Vector(LATENT_SIZE)) + .unwrap(), + neura_layer!("softmax") + ]; + let test_data = test_images + .clone() + .zip(test_labels.clone()) + .collect::>(); + + let trainer = NeuraBatchedTrainer::with_epochs(0.03, 20, 128, TRAIN_SIZE); + + plot_losses( + trainer.train( + &NeuraBackprop::new(Euclidean), + &mut network, + cycle_shuffling(train_images.clone().zip(train_labels), rand::thread_rng()), + &test_data, + ), + 128, + 48, + ); + + let mut correct = 0; + for (test_image, test_label) in test_images.zip(test_labels) { + let guess = network.eval(&test_image); + let guess = argmax(guess.as_slice()); + let actual = argmax(test_label.as_slice()); + + if guess == actual { + correct += 1; + } + } + + println!(""); + println!( + "{} correct out of {}: {:.2}%", + correct, + TEST_SIZE, + (correct as f32 / TEST_SIZE as f32) * 100.0 + ); +} + +fn one_hot(value: usize, categories: usize) -> DVector { + let mut res = DVector::from_element(categories, 0.0); + if value < categories { + res[value] = 1.0; + } + res +} diff --git a/examples/mnist-diffusion.rs b/examples/mnist-diffusion.rs new file mode 100644 index 0000000..3b7b062 --- /dev/null +++ b/examples/mnist-diffusion.rs @@ -0,0 +1,213 @@ +use nalgebra::DVector; +use rand::Rng; +use rust_mnist::Mnist; +use std::io::Write; + +use neuramethyst::{ + cycle_shuffling, + derivable::{ + activation::{Logistic, Relu, Swish}, + loss::Euclidean, + regularize::NeuraL2, + }, + plot_losses, + prelude::*, +}; + +const TRAIN_SIZE: usize = 50000; +const TEST_SIZE: usize = 1000; +const WIDTH: usize = 28; +const HEIGHT: usize = 28; +const REG_RATE: f32 = 0.003; +const EPOCHS: usize = 80; + +// const BASE_NOISE: f32 = 0.05; +const NOISE_AMOUNT: f32 = 0.5; +const SHIFT_AMOUNT: i32 = 9; + +pub fn main() { + let Mnist { + train_data: train_images, + train_labels, + test_data: test_images, + test_labels, + .. + } = Mnist::new("data/"); + + let train_images = train_images + .into_iter() + .map(|raw| { + DVector::from_iterator(WIDTH * HEIGHT, raw.into_iter().map(|x| x as f32 / 255.0)) + }) + .take(TRAIN_SIZE); + let train_labels = train_labels + .into_iter() + .map(|x| one_hot(x as usize, 10)) + .take(TRAIN_SIZE); + + let test_images = test_images + .into_iter() + .map(|raw| { + DVector::from_iterator(WIDTH * HEIGHT, raw.into_iter().map(|x| x as f32 / 255.0)) + }) + .take(TEST_SIZE); + let test_labels = test_labels + .into_iter() + .map(|x| one_hot(x as usize, 10)) + .take(TEST_SIZE); + + let test_data: Vec<_> = augment_data(test_images.zip(test_labels)).collect(); + + let mut network = neura_residual![ + <= 0, 1; + neura_layer!("isolate", WIDTH * HEIGHT, WIDTH * HEIGHT + 10) => 1, 3, 5, 7, 9, 10; + neura_layer!("isolate", 0, WIDTH * HEIGHT) => 0, 1, 3; + neura_layer!("dense", 100).regularization(NeuraL2(REG_RATE)).activation(Swish(Logistic)) => 0, 2; + neura_layer!("dropout", 0.5); + neura_layer!("dense", 50).regularization(NeuraL2(REG_RATE)).activation(Swish(Logistic)) => 0, 2, 4; + neura_layer!("dropout", 0.5); + neura_layer!("dense", 50).regularization(NeuraL2(REG_RATE)).activation(Swish(Logistic)) => 0, 2; + neura_layer!("dropout", 0.33); + neura_layer!("dense", 25).regularization(NeuraL2(REG_RATE)).activation(Swish(Logistic)) => 0, 2; + neura_layer!("dropout", 0.33); + neura_layer!("dense", 25).regularization(NeuraL2(REG_RATE)).activation(Swish(Logistic)); + // neura_layer!("dropout", 0.33); + neura_layer!("dense", WIDTH * HEIGHT).activation(Relu); + ] + .construct(NeuraShape::Vector(WIDTH * HEIGHT + 10)) + .unwrap(); + + let trainer = NeuraBatchedTrainer::with_epochs(0.03, EPOCHS, 512, TRAIN_SIZE); + // trainer.log_iterations = 1; + let train_data = augment_data(cycle_shuffling( + train_images.clone().zip(train_labels.clone()), + rand::thread_rng(), + )); + + let losses = trainer.train( + &NeuraBackprop::new(Euclidean), + &mut network, + train_data, + &test_data, + ); + plot_losses(losses, 128, 48); + + loop { + let mut image = uniform_vector(WIDTH * HEIGHT + 10); + let mut buffer = String::new(); + print!("> "); + std::io::stdout().flush().unwrap(); + if let Err(_) = std::io::stdin().read_line(&mut buffer) { + break; + } + + for i in 0..10 { + image[WIDTH * HEIGHT + i] = buffer + .chars() + .any(|c| c == char::from_digit(i as u32, 10).unwrap()) + as u8 as f32; + } + + for _iter in 0..5 { + let new_image = network.eval(&image); + + neuramethyst::draw_neuron_activation( + |[x, y]| { + let x = ((x + 1.0) / 2.0 * WIDTH as f32) as usize; + let y = ((y + 1.0) / 2.0 * HEIGHT as f32) as usize; + + let index = x + y * WIDTH; + + vec![new_image[index]] + }, + 1.0, + WIDTH as u32, + HEIGHT as u32, + ); + + for i in 0..(WIDTH * HEIGHT) { + image[i] = new_image[i] * 0.6 + image[i] * 0.3; + } + + std::thread::sleep(std::time::Duration::new(0, 100_000_000)); + } + } +} + +fn uniform_vector(length: usize) -> DVector { + let mut res = DVector::from_element(length, 0.0); + let mut rng = rand::thread_rng(); + + for i in 0..length { + res[i] = rng.gen(); + } + + res +} + +fn one_hot(value: usize, categories: usize) -> DVector { + let mut res = DVector::from_element(categories, 0.0); + if value < categories { + res[value] = 1.0; + } + res +} + +fn add_noise(mut image: DVector, rng: &mut impl Rng, amount: f32) -> DVector { + if amount <= 0.0 { + return image; + } + + let uniform = rand::distributions::Uniform::new(0.0, amount); + + for i in 0..image.len() { + let x = rng.sample(uniform); + image[i] = image[i] * (1.0 - x) + (1.0 - image[i]) * x; + } + + image +} + +fn shift(image: &DVector, dx: i32, dy: i32) -> DVector { + let mut res = DVector::from_element(image.len(), 0.0); + let width = WIDTH as i32; + let height = HEIGHT as i32; + + for y in 0..height { + for x in 0..width { + let x2 = x + dx; + let y2 = y + dy; + if y2 < 0 || y2 >= height || x2 < 0 || x2 >= width { + continue; + } + res[(y2 * width + x2) as usize] = image[(y * width + x) as usize]; + } + } + + res +} + +fn augment_data( + iter: impl Iterator, DVector)>, +) -> impl Iterator, DVector)> { + let mut rng = rand::thread_rng(); + iter.map(move |(image, label)| { + let noise_amount = rng.gen_range(0.05..NOISE_AMOUNT); + let base_image = shift( + &image, + rng.gen_range(-SHIFT_AMOUNT..SHIFT_AMOUNT), + rng.gen_range(-SHIFT_AMOUNT..SHIFT_AMOUNT), + ) * rng.gen_range(0.6..1.0); + // let base_image = add_noise(base_image, &mut rng, base_noise); + + let noisy_image = add_noise(base_image.clone(), &mut rng, noise_amount); + + ( + DVector::from_iterator( + WIDTH * HEIGHT + 10, + noisy_image.iter().copied().chain(label.iter().copied()), + ), + image, + ) + }) +} diff --git a/src/layer/isolate.rs b/src/layer/isolate.rs new file mode 100644 index 0000000..049bec5 --- /dev/null +++ b/src/layer/isolate.rs @@ -0,0 +1,161 @@ +use nalgebra::{DVector, Scalar}; + +use super::*; + +/// **Class invariant:** start and end are +#[derive(Clone, Debug)] +pub struct NeuraIsolateLayer { + start: NeuraShape, + end: NeuraShape, +} + +#[derive(Clone, Debug)] +pub enum NeuraIsolateLayerErr { + Incompatible { + start: NeuraShape, + end: NeuraShape, + input_shape: NeuraShape, + }, + OutOfBound { + start: NeuraShape, + end: NeuraShape, + input_shape: NeuraShape, + }, + OutOfOrder { + start: NeuraShape, + end: NeuraShape, + }, +} + +impl NeuraIsolateLayer { + pub fn new>(start: T, end: T) -> Option { + let start = start.into(); + let end = end.into(); + + if start.is_compatible(end) { + Some(Self { start, end }) + } else { + None + } + } +} + +impl NeuraShapedLayer for NeuraIsolateLayer { + fn output_shape(&self) -> NeuraShape { + self.end.sub(self.start).unwrap_or_else(|| unreachable!()) + } +} + +impl NeuraPartialLayer for NeuraIsolateLayer { + type Constructed = NeuraIsolateLayer; + type Err = NeuraIsolateLayerErr; + + fn construct(self, input_shape: NeuraShape) -> Result { + use NeuraShape::*; + let start = self.start; + let end = self.end; + + match (input_shape, start, end) { + (Vector(xi), Vector(xs), Vector(xe)) => { + if xs >= xe { + return Err(NeuraIsolateLayerErr::OutOfOrder { start, end }); + } + + if xs >= xi || xe > xi { + return Err(NeuraIsolateLayerErr::OutOfBound { + start, + end, + input_shape, + }); + } + + Ok(self) + } + + (Matrix(_xi, _yi), Matrix(_xs, _ys), Matrix(_xe, _ye)) => unimplemented!(), + (Tensor(_xi, _yi, _zi), Tensor(_xs, _ys, _zs), Tensor(_xe, _ye, _ze)) => { + unimplemented!() + } + + _ => Err(NeuraIsolateLayerErr::Incompatible { + start, + end, + input_shape, + }), + } + } +} + +impl NeuraLayer> for NeuraIsolateLayer { + type Output = DVector; + + fn eval(&self, input: &DVector) -> Self::Output { + let (NeuraShape::Vector(start), NeuraShape::Vector(end)) = (self.start, self.end) else { + panic!("NeuraIsolateLayer expected a value of dimension {}, got a vector", self.start.dims()); + }; + + DVector::from_iterator(end - start, input.iter().cloned().skip(start).take(end)) + } +} + +impl NeuraTrainableLayerBase for NeuraIsolateLayer { + type Gradient = (); + + #[inline(always)] + fn default_gradient(&self) -> Self::Gradient { + () + } + + #[inline(always)] + fn apply_gradient(&mut self, _gradient: &Self::Gradient) { + // Noop + } +} + +impl NeuraTrainableLayerEval> for NeuraIsolateLayer { + type IntermediaryRepr = (); + + fn eval_training(&self, input: &DVector) -> (Self::Output, Self::IntermediaryRepr) { + (self.eval(input), ()) + } +} + +impl NeuraTrainableLayerSelf for NeuraIsolateLayer +where + Self: NeuraTrainableLayerEval, +{ + #[inline(always)] + fn regularize_layer(&self) -> Self::Gradient { + () + } + + #[inline(always)] + fn get_gradient( + &self, + _input: &Input, + _intermediary: &Self::IntermediaryRepr, + _epsilon: &Self::Output, + ) -> Self::Gradient { + () + } +} + +impl NeuraTrainableLayerBackprop> for NeuraIsolateLayer { + fn backprop_layer( + &self, + input: &DVector, + _intermediary: &Self::IntermediaryRepr, + epsilon: &Self::Output, + ) -> DVector { + let mut result = DVector::from_element(input.len(), F::default()); + let NeuraShape::Vector(start) = self.start else { + unreachable!(); + }; + + for i in 0..epsilon.len() { + result[start + i] = epsilon[i].clone(); + } + + result + } +} diff --git a/src/layer/mod.rs b/src/layer/mod.rs index 43f38d7..9dd6547 100644 --- a/src/layer/mod.rs +++ b/src/layer/mod.rs @@ -4,6 +4,7 @@ use self::lock::NeuraLockLayer; pub mod dense; pub mod dropout; +pub mod isolate; pub mod lock; pub mod normalize; pub mod softmax; @@ -23,6 +24,55 @@ impl NeuraShape { NeuraShape::Tensor(rows, columns, channels) => rows * columns * channels, } } + + pub fn sub(&self, other: NeuraShape) -> Option { + use NeuraShape::*; + + Some(match (other, self) { + (Vector(x1), Vector(x2)) => Vector(x2 - x1), + (Matrix(x1, y1), Matrix(x2, y2)) => Matrix(x2 - x1, y2 - y1), + (Tensor(x1, y1, z1), Tensor(x2, y2, z2)) => Tensor(x2 - x1, y2 - y1, z2 - z1), + + _ => return None, + }) + } + + pub fn is_compatible(&self, other: NeuraShape) -> bool { + use NeuraShape::*; + + matches!( + (self, other), + (Vector(_), Vector(_)) + | (Matrix(_, _), Matrix(_, _)) + | (Tensor(_, _, _), Tensor(_, _, _)) + ) + } + + pub fn dims(&self) -> usize { + match self { + NeuraShape::Vector(_) => 1, + NeuraShape::Matrix(_, _) => 2, + NeuraShape::Tensor(_, _, _) => 3, + } + } +} + +impl From for NeuraShape { + fn from(x: usize) -> Self { + NeuraShape::Vector(x) + } +} + +impl From<(usize, usize)> for NeuraShape { + fn from((x, y): (usize, usize)) -> Self { + NeuraShape::Matrix(x, y) + } +} + +impl From<(usize, usize, usize)> for NeuraShape { + fn from((x, y, z): (usize, usize, usize)) -> Self { + NeuraShape::Tensor(x, y, z) + } } pub trait NeuraLayer { @@ -217,4 +267,8 @@ macro_rules! neura_layer { ( "normalize" ) => { $crate::layer::normalize::NeuraNormalizeLayer::new() }; + + ( "isolate", $start:expr, $end:expr ) => { + $crate::layer::isolate::NeuraIsolateLayer::new($start, $end).unwrap() + }; } diff --git a/src/lib.rs b/src/lib.rs index 0eeb5b2..201c9c9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,6 +10,9 @@ mod utils; // TODO: move to a different file pub use utils::{argmax, cycle_shuffling, one_hot, plot_losses}; +#[cfg(feature = "visualization")] +pub use utils::draw_neuron_activation; + /// Common traits and structs that are useful to use this library. /// All of these traits are prefixed with the word "neura" in some way, /// so there should not be any conflicts when doing a wildcard import of `prelude`. diff --git a/src/network/residual/wrapper.rs b/src/network/residual/wrapper.rs index a0df715..f3cdc7e 100644 --- a/src/network/residual/wrapper.rs +++ b/src/network/residual/wrapper.rs @@ -1,6 +1,6 @@ use std::borrow::Cow; -use crate::{network::*, utils::unwrap_or_clone}; +use crate::network::*; use super::*; diff --git a/src/utils.rs b/src/utils.rs index 9b2191d..259c65e 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -97,6 +97,7 @@ pub(crate) fn uniform_vector(length: usize) -> nalgebra::DVector { DVector::from_fn(length, |_, _| -> f64 { rng.gen() }) } +#[deprecated] pub fn one_hot(value: usize) -> NeuraVector { let mut res = NeuraVector::default(); if value < N { @@ -154,3 +155,41 @@ pub(crate) fn unwrap_or_clone(value: std::rc::Rc) -> T { // TODO: replace with Rc::unwrap_or_clone once https://github.com/rust-lang/rust/issues/93610 is closed std::rc::Rc::try_unwrap(value).unwrap_or_else(|value| (*value).clone()) } + +#[cfg(feature = "visualization")] +pub fn draw_neuron_activation Vec>( + callback: F, + scale: f32, + width: u32, + height: u32, +) { + use viuer::Config; + + let mut image = image::RgbImage::new(width, height); + + fn sigmoid(x: f32) -> f32 { + 1.9 / (1.0 + (-x * 3.0).exp()) - 0.9 + } + + for y in 0..height { + let y2 = 2.0 * y as f32 / height as f32 - 1.0; + for x in 0..width { + let x2 = 2.0 * x as f32 / width as f32 - 1.0; + let activation = callback([x2 * scale, y2 * scale]); + let r = (sigmoid(activation.get(0).copied().unwrap_or(-1.0)) * 255.0).floor() as u8; + let g = (sigmoid(activation.get(1).copied().unwrap_or(-1.0)) * 255.0).floor() as u8; + let b = (sigmoid(activation.get(2).copied().unwrap_or(-1.0)) * 255.0).floor() as u8; + + *image.get_pixel_mut(x, y) = image::Rgb([r, g, b]); + } + } + + let config = Config { + use_kitty: false, + truecolor: true, + // absolute_offset: false, + ..Default::default() + }; + + viuer::print(&image::DynamicImage::ImageRgb8(image), &config).unwrap(); +}