use gtensor as gt;
use rand::Rng;
const N: usize = 10;
fn main() {
let train = load_dataset();
let test = load_dataset();
draw_dataset(&train, "examples/class-train-dataset.bmp");
draw_dataset(&test, "examples/class-test-dataset.bmp");
let mut tape = build_tape();
tape.set_batch_size(N);
let mut error = gt::Tensor::from_fill([N], 0.0);
let mut grad = gt::Tensor::from_fill([N,1], 0.0);
for epoch in 0..50 {
let mut loss = 0.0;
for (feature, label) in train.iter_batched(N) {
{
let prediction = tape.forward(feature);
gt::math::loss::mse(
label,
prediction.slice(),
&mut error,
&mut grad,
);
}
error.iter().for_each(|x| loss += *x);
tape.backward(grad.slice())
}
println!("Epoch: {epoch}, loss: {}", loss / 200.)
}
let mut loss = 0.0;
for (feature, label) in test.iter_batched(N) {
let prediction = tape.forward(feature);
gt::math::loss::mse(
label,
prediction.slice(),
&mut error,
&mut grad,
);
error.iter().for_each(|x| loss += *x);
}
println!("Test Loss: {}", loss / 200.);
draw_prediction(&mut tape, &train, "examples/class-prediction.bmp");
tape.save("examples/class").unwrap();
}
fn load_dataset() -> gt::Dataset {
let mut data = gt::Dataset::new([2], [1]);
let mut rng = rand::thread_rng();
for _ in 0..200 {
let x = rng.gen_range(-1.0..1.0);
let y = rng.gen_range(-1.0..1.0);
let d = f32::sqrt((x*x)+(y*y));
let d = if d > 0.6 { 1.0 } else { 0.0 };
data.load_feature(&[x,y], &[d]);
}
data
}
fn build_tape() -> gt::Tape {
let mut tape = gt::Tape::builder();
tape.opt = gt::opt::momentum(0.04, 0.9);
tape.init = gt::init::normal(0.5, 1.0);
let x = tape.input([2]);
let w = tape.parameter([2,4]);
let b = tape.parameter([4]);
let x = gt::op::matmul(x, w);
let x = gt::op::axis_add(x, b, 'C');
let x = gt::op::tanh(x);
let w = tape.parameter([4,4]);
let b = tape.parameter([4]);
let x = gt::op::matmul(x,w);
let x = gt::op::axis_add(x, b, 'C');
let x = gt::op::tanh(x);
let w = tape.parameter([4,1]);
let b = tape.parameter([1]);
let x = gt::op::matmul(x, w);
let x = gt::op::axis_add(x, b, 'C');
let _ = gt::op::tanh(x);
tape.finish()
}
pub fn draw_dataset(dataset: >::Dataset, name: &str) {
let mut img = bmp::Image::new(200,200);
for x in 0..200 {
for y in 0..200 {
img.set_pixel(x, y, bmp::Pixel::new(255,255,255))
}
}
for (feature, label) in dataset.iter_batched(1) {
let color =
if label[0] < 0.5 {
bmp::Pixel::new(60, 165, 255)
} else {
bmp::Pixel::new(255, 165, 0)
};
let x = ((feature[0] + 1.0) * 100.) as u32;
let y = ((feature[1] + 1.0) * 100.) as u32;
img.set_pixel(x, y, color)
}
img.save(name).unwrap()
}
pub fn draw_prediction(tape: &mut gt::Tape, dataset: >::Dataset, name: &str) {
let mut img = bmp::Image::new(200,200);
tape.set_batch_size(1);
for x in 0..200 {
for y in 0..200 {
let xf = (x as f32 / 100.) - 1.;
let yf = (y as f32 / 100.) - 1.;
let input = gt::Tensor::from_slice([1,2], &[xf,yf]);
let prediction = tape.forward(input.slice());
let color =
if prediction[0] < 0.5 {
bmp::Pixel::new(60, 165, 255)
} else {
bmp::Pixel::new(255, 165, 0)
};
img.set_pixel(x, y, color)
}
}
for (feature, label) in dataset.iter_batched(1) {
let color =
if label[0] < 0.5 {
bmp::Pixel::new(30, 135, 225)
} else {
bmp::Pixel::new(225, 135, 0)
};
let x = ((feature[0] + 1.0) * 100.) as u32;
let y = ((feature[1] + 1.0) * 100.) as u32;
img.set_pixel(x, y, color)
}
img.save(name).unwrap();
}