topos 0.12.0

An autodiff compiler stack in 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
112
113
114
115
116
117

//! 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, 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::{Activation, Mlp, Module, Shape, Tape, Tensor, 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 tape: Tape<f32> = Tape::new();
    let mlp = Mlp::new(
        &tape,
        &[1, HIDDEN_LEN, 1],
        Activation::Tanh,
        init::xavier(7),
    );

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

    let (input, predicted, loss) = (input.symbol(), predicted.symbol(), loss.symbol());
    let network = tape.into_network();
    let mut parameters = network.parameters();

    let learning_rate = Tensor::new([], [0.001]);
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..STEP_COUNT {
        let run = network.forward(&parameters, []);
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        if step % (STEP_COUNT / 5) == 0 {
            println!("step {step:4}: loss = {batch_loss:.4}");
        }
        let gradients = run.backward(loss).parameters(&parameters);
        parameters = parameters.step(&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(
        &parameters,
        [(input, Tensor::new([SAMPLE_LEN, 1], grid.clone()))],
    );
    let fit = run.of(predicted).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())
    );
}