From 636447a18ad5ce4d88d858c6e8d0d521c4ec7dee Mon Sep 17 00:00:00 2001 From: Vishal Patil Date: Thu, 16 Jan 2025 21:06:51 -0500 Subject: [PATCH] Function approximator --- neural-networks-101/Cargo.toml | 16 ++++ neural-networks-101/src/main.rs | 151 ++++++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+) create mode 100644 neural-networks-101/Cargo.toml create mode 100644 neural-networks-101/src/main.rs diff --git a/neural-networks-101/Cargo.toml b/neural-networks-101/Cargo.toml new file mode 100644 index 0000000..df2a8e8 --- /dev/null +++ b/neural-networks-101/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "neural-networks-101" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +anyhow = "1.0.40" +csv = "1.1.6" +clap = { version = "4.5.1", features = ["derive"] } +rand = "0.8.5" +candle-core = { git = "https://github.com/huggingface/candle.git", tag = "0.6.0", features = [ + "cuda", +] } +candle-nn = { git = "https://github.com/huggingface/candle.git", tag = "0.6.0" } diff --git a/neural-networks-101/src/main.rs b/neural-networks-101/src/main.rs new file mode 100644 index 0000000..f0dc2f7 --- /dev/null +++ b/neural-networks-101/src/main.rs @@ -0,0 +1,151 @@ +use anyhow::Result; +use candle_core::{DType, Device, Tensor, D}; +use candle_nn::{loss, ops, Linear, Module, Optimizer, VarBuilder, VarMap}; +use clap::Parser; +use rand::Rng; +use std::f64::consts::PI; +use std::rc::Rc; + +const INPUT_COUNT: usize = 3; +const LAYER1_COUNT: usize = 100; +const LAYER2_COUNT: usize = 100; +const OUTPUT_COUNT: usize = 1; + +struct Dataset { + pub training_data: Tensor, + pub training_labels: Tensor, + pub test_data: Tensor, + pub test_labels: Tensor, +} + +fn func(x1: f32, x2: f32, x3: f32) -> f32 { + 2.0 * (x1.abs() + 1.0).ln() + (3.0 * x2).sin() * (-0.5 * x3).exp() + 0.5 * x1 * x3 +} + +fn generate_nonlinear_data(n_samples: usize) -> (Vec>, Vec) { + let mut rng = rand::thread_rng(); + let mut x_values = Vec::with_capacity(n_samples); + let mut y_values = Vec::with_capacity(n_samples); + + for _ in 0..n_samples { + let x1 = rng.gen_range(-5.0..=5.0) as f32; + let x2 = rng.gen_range(-PI..=PI) as f32; + let x3 = rng.gen_range(-3.0..=3.0) as f32; + let noise = rng.gen_range(-0.5..=0.5) as f32; + let y = (func(x1, x2, x3) + noise) as f32; + x_values.push(vec![x1, x2, x3]); + y_values.push(y); + } + + (x_values, y_values) +} + +fn load_tensors(samples: u32, device: &Device) -> Result<(Tensor, Tensor)> { + let (x_values, y_values) = generate_nonlinear_data(samples as usize); + + let x_values = x_values.into_iter().flatten().collect::>(); + let x_values = x_values.as_slice(); + let x_tensor = Tensor::from_slice(x_values, &[samples as usize, 3], &device)?; + + let y_values = y_values.into_iter().collect::>(); + let y_values = y_values.as_slice(); + let y_tensor = Tensor::from_slice(y_values, &[samples as usize], &device)?; + + Ok((x_tensor, y_tensor)) +} + +fn load_dataset(device: &Device) -> Result { + let (training_data, training_labels) = load_tensors(5000, device)?; + let (test_data, test_labels) = load_tensors(2000, device)?; + + Ok(Dataset { + training_data, + training_labels, + test_data, + test_labels, + }) +} + +struct Mlp { + ln1: Linear, + ln2: Linear, + ln3: Linear, +} + +impl Mlp { + fn new(vs: VarBuilder) -> Result { + let ln1 = candle_nn::linear(INPUT_COUNT, LAYER1_COUNT, vs.pp("ln1"))?; + let ln2 = candle_nn::linear(LAYER1_COUNT, LAYER2_COUNT, vs.pp("ln2"))?; + let ln3 = candle_nn::linear(LAYER2_COUNT, OUTPUT_COUNT, vs.pp("ln3"))?; + Ok(Self { ln1, ln2, ln3 }) + } + + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.ln1.forward(xs)?; + let xs = xs.tanh()?; + let xs = self.ln2.forward(&xs)?; + let xs = xs.tanh()?; + Ok(self.ln3.forward(&xs)?) + } +} + +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +struct Args { + // Print the Cost and Loss at each epoch + #[arg(long, default_value_t = false)] + progress: bool, + + // The learning rate + #[arg(long, default_value = "0.01")] + learning_rate: f64, + + // The regularization parameter + #[arg(long, default_value = "0.01")] + regularization: f32, + + // The number of epochs + #[arg(long, default_value = "5000")] + epochs: i32, +} + +fn main() -> Result<()> { + let args = Args::parse(); + let device = Rc::new(Device::cuda_if_available(0)?); + let dataset = load_dataset(&device)?; + + let varmap = VarMap::new(); + let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device); + let model = Mlp::new(vs)?; + let mut sgd = candle_nn::SGD::new(varmap.all_vars(), args.learning_rate)?; + + let test_images = dataset.test_data.to_device(&device)?; + let test_labels = dataset + .test_labels + .to_dtype(DType::U32)? + .to_device(&device)?; + + for epoch in 1..args.epochs { + let logits = model.forward(&dataset.training_data)?; + let loss = loss::mse(&logits.squeeze(1)?, &dataset.training_labels)?; + sgd.backward_step(&loss)?; + + let test_logits = model.forward(&test_images)?; + let sum_ok = test_logits + .argmax(D::Minus1)? + .eq(&test_labels)? + .to_dtype(DType::F32)? + .sum_all()? + .to_scalar::()?; + let test_accuracy = sum_ok / test_labels.dims1()? as f32; + if args.progress && epoch % 100 == 0 { + println!( + "{epoch:4} train loss: {:8.5} test acc: {:5.2}%", + loss.to_scalar::()?, + 100. * test_accuracy + ); + } + } + + Ok(()) +}