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
118
119

//! Trains a linear model `w * x + b` with gradient descent, exercising one
//! shared network across threads.
//!
//! Two things run in parallel here. First, a single run of the shared
//! network feeds concurrent backward sweeps, one per target: runs are
//! per-thread state, the network is immutable. Second, several training
//! runs proceed simultaneously over one spec — each learning rate owns
//! its own `Parameters` state, and cloning state is all a what-if costs.
//!
//! Run with: `cargo run --example gradient_descent`

use std::time::Instant;

use rayon::prelude::*;

use malevich::{Frame, Line, Plot};
use topos::model::{Detach, Tape, Tensor};

fn main() {
    // The whole recording is one closure; its return value is the
    // set of names that leave the tape, detached in one `detach` call.
    let (network, (w, b, loss, sample_losses)) = Tape::record(|tape| {
        // Learnable parameters, starting from zero.
        let w = tape.parameter(0.0_f64);
        let b = tape.parameter(0.0);

        // Training data for the target line `y = 2 * x + 1`, recorded
        // as plain leaves. Each sample's squared error is kept as a
        // separate target; the total loss is their sum.
        let samples = [(1.0, 3.0), (2.0, 5.0), (3.0, 7.0)];
        let mut sample_losses = Vec::new();
        for (x, y) in samples {
            let x = tape.leaf(x);
            let y = tape.leaf(y);
            let error = w * x + b - y;
            sample_losses.push(error * error);
        }
        // Values are `Copy`, so the per-sample losses fold into a total
        // loss with a plain reduce.
        let loss = sample_losses
            .iter()
            .copied()
            .reduce(|total, squared| total + squared)
            .expect("at least one sample");
        (w, b, loss, sample_losses).detach()
    });
    let parameters = network.parameters();

    // One run feeds many backward sweeps: each rayon thread
    // differentiates the same shared run for its own target.
    let run = network.forward(&parameters, []);
    let per_sample: Vec<f64> = sample_losses
        .par_iter()
        .map(|&sample_loss| {
            let gradients = run.backward(sample_loss);
            gradients.of(w).scalar()
        })
        .collect();
    let total_gradient = run.backward(loss).of(w).scalar();
    println!("per-sample d/dw, computed on separate threads: {per_sample:?}");
    println!(
        "their sum {} equals the total-loss d/dw {} by linearity",
        per_sample.iter().sum::<f64>(),
        total_gradient
    );

    // Parallel training: each learning rate clones the parameter state
    // and descends independently over the one shared spec, keeping its
    // whole loss history for the chart.
    let learning_rates = [0.005, 0.02, 0.05];
    let training = Instant::now();
    let runs: Vec<(f64, Vec<f64>, f64, f64)> = learning_rates
        .par_iter()
        .map(|&learning_rate| {
            let mut parameters = parameters.clone();
            let mut losses = Vec::with_capacity(501);
            for _ in 0..500 {
                let run = network.forward(&parameters, []);
                losses.push(run.of(loss).scalar());
                let gradients = run.backward(loss).parameters(&parameters);
                parameters = parameters.step(&gradients, |parameter, gradient| {
                    parameter.clone() - gradient.clone() * Tensor::from(learning_rate)
                });
            }
            let run = network.forward(&parameters, []);
            losses.push(run.of(loss).scalar());
            (
                learning_rate,
                losses,
                parameters.of(w).scalar(),
                parameters.of(b).scalar(),
            )
        })
        .collect();

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

    println!("parallel training on cloned states (target: w = 2, b = 1):");
    for (learning_rate, losses, w, b) in &runs {
        let final_loss = losses.last().expect("every run records its final loss");
        println!("  lr = {learning_rate:5.3}: loss = {final_loss:.6}, w = {w:.3}, b = {b:.3}");
    }

    // The same three descents as curves: on a log scale a constant
    // convergence rate is a straight line, so the slope is the rate.
    let mut plot = Plot::new()
        .title("gradient descent per learning rate")
        .x_label("step")
        .y_label("loss")
        .log_y();
    for (learning_rate, losses, ..) in &runs {
        plot = plot.layer(Line::y(&losses[..]).label(format!("lr = {learning_rate}")));
    }
    println!("{}", plot.render_best(&Frame::detect()));
}