topos 0.11.0

A tiny autograd engine for the GPU-poor, written in idiomatic Rust.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

//! Fits a one-hidden-layer MLP to a noisy sine wave and draws the fit:
//! the sampled points, the curve they came from, and the curve the
//! model learned, in one chart.
//!
//! The graph is recorded once over a `[SAMPLE_LEN, 1]` input whose
//! default feed is the training set, so training is a plain `forward`
//! per step; charting feeds an evenly spaced grid through the same
//! expression with `forward_with`, so the fit line is the model itself,
//! not an interpolation of its training outputs.
//!
//! Run with: `cargo run --release --example regression`

use std::time::Instant;

use malevich::{Frame, Line, Plot, Points};
use topos::{Mlp, Network, Shape, Tensor, Tensorial, init};

/// How many noisy samples the model trains on; the chart grid reuses
/// the count so the recorded expression serves both.
const SAMPLE_LEN: usize = 96;

/// The half-width of the sampled domain, centered on zero.
const DOMAIN: f64 = 3.0;

/// The standard deviation of the noise added to the sampled curve.
const NOISE_DEVIATION: f64 = 0.1;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 16;

/// How many full-batch gradient descent steps the fit takes.
const STEP_COUNT: usize = 3000;

fn main() {
    // The dataset: x uniform over the domain and y = sin x plus noise,
    // both drawn from the same seeded initializers the parameters use.
    let features: Tensor<f32> = init::uniform(3, DOMAIN)(&Shape::new([SAMPLE_LEN, 1]));
    let noise: Tensor<f32> = init::normal(5, NOISE_DEVIATION)(&Shape::new([SAMPLE_LEN, 1]));
    let feature_values = features.to_vec();
    let target_values: Vec<f32> = feature_values
        .iter()
        .zip(noise.to_vec())
        .map(|(&x, noise)| x.sin() + noise)
        .collect();

    let network: Network<Tensor<f32>> = Network::new();
    let mlp = Mlp::new(&network, &[1, HIDDEN_LEN, 1], init::xavier(7));

    let input = network.input(features);
    let expected = network.input(Tensor::new([SAMPLE_LEN, 1], target_values.clone()));
    let predicted = mlp.express(&network, input);
    let error = predicted - expected;
    let loss = (error * error).sum();

    let input_symbol = input.symbol();
    let predicted_symbol = predicted.symbol();
    let loss_symbol = loss.symbol();

    let learning_rate = Tensor::new([], [0.001]);
    let mut network = network;
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..STEP_COUNT {
        let loss_value = network.resolve(loss_symbol);
        let run = network.forward();
        let batch_loss = run.of(loss_value).to_vec()[0];
        losses.push(batch_loss);
        if step % (STEP_COUNT / 5) == 0 {
            println!("step {step:4}: loss = {batch_loss:.4}");
        }
        let gradients = run.backward(loss_value);
        network = network.update(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    println!(
        "{}",
        Plot::new()
            .layer(Line::y(&losses[..]).label("full batch"))
            .title("regression training")
            .x_label("step")
            .y_label("sum of squared errors")
            .render_best(&Frame::detect())
    );

    // The fit, read over an evenly spaced grid fed through the trained
    // expression in place of the training samples.
    let grid: Vec<f32> = (0..SAMPLE_LEN)
        .map(|index| (((index as f64 / (SAMPLE_LEN - 1) as f64) * 2.0 - 1.0) * DOMAIN) as f32)
        .collect();
    let run = network.forward_with([(input_symbol, Tensor::new([SAMPLE_LEN, 1], grid.clone()))]);
    let fit = run.of(network.resolve(predicted_symbol)).to_vec();

    println!(
        "{}",
        Plot::new()
            .layer(Points::xy(&feature_values[..], &target_values[..]).label("samples"))
            .layer(Line::function(-DOMAIN..DOMAIN, f64::sin).label("sin x"))
            .layer(Line::xy(&grid[..], &fit[..]).label("fit"))
            .title("a tanh mlp fit to noisy sin x")
            .x_label("x")
            .render_best(&Frame::detect())
    );
}