ex_ml_regression.rs

1use oa::{Engine, Matrix, ml};
2
3fn main() -> oa::Result<()> {
4	let engine = Engine::new()?;
5
6	let x = Matrix::from_f32(&engine, [5, 1], &[-2.0, -1.0, 0.0, 1.0, 2.0])?;
7	let target = Matrix::from_f32(&engine, [5, 1], &[-3.0, -1.0, 1.0, 3.0, 5.0])?;
8
9	let model = ml::nn::Linear::with_seed(&engine, 1, 1, 0)?;
10	let mut optimizer = ml::Sgd::new(model.parameters(), 0.05, 0.0, 0.0)?;
11
12	let mut initial_loss = 0.0_f32;
13	let mut loss_matrix = Matrix::from_f32(&engine, [1], &[0.0])?;
14
15	for step in 0..80 {
16		optimizer.zero_grad();
17		let tape = ml::GradientTape::new();
18		let prediction = model.forward(&x)?;
19		loss_matrix = ml::loss::mse(&prediction, &target)?;
20		tape.backward(&loss_matrix)?;
21		drop(tape);
22		optimizer.step()?;
23
24		if step == 0 {
25			initial_loss = loss_matrix.read_f32()?[0];
26		}
27	}
28
29	let final_loss = loss_matrix.read_f32()?[0];
30	let values = model.forward(&x)?.read_f32()?;
31
32	assert!(final_loss < initial_loss * 0.01);
33	assert!(
34		values
35			.iter()
36			.zip([-3.0_f32, -1.0, 1.0, 3.0, 5.0])
37			.all(|(&a, b)| (a - b).abs() < 0.1)
38	);
39
40	println!("loss: {initial_loss:.6} -> {final_loss:.6}");
41	println!("{values:?}");
42	Ok(())
43}
44