3. Neural Networks

RustyML holds a small deep-learning framework, written in pure Rust. Every tensor of that framework is a Tensor, an alias for ndarray::ArrayD<f32>. A tensor holds single-precision values and has a dynamic rank. The framework runs on the CPU and keeps no autograd tape. Each layer writes its own forward pass and its own backward pass, and hands its parameters to the optimizer as a flat view. The framework is therefore deterministic and easy to debug.

A layer computes, and it remembers nothing. forward and backward both take &self and a &mut Ctx. The context carries the training flag, the caches of the pass, the parameter gradients, and the non-trainable state that a training pass changes. A built model is therefore Send and Sync, and several threads can run inference against 1 model. 3 things separate this framework from most others. It keeps a strict f32 precision, it reads the input shape 1 time at a model’s build, and every constructor returns a Result.

The crate holds 2 model types, and both drive the same layers. Sequential is a chain. Every layer of a chain takes 1 input, and that input is the output of the layer before it. Graph is a directed graph, and a node of a graph is 1 call of 1 layer on the outputs of other nodes. A graph model can therefore hold several inlets, several outlets, a residual connection, or 1 layer that several nodes share. GraphBuilder collects the nodes, and its build refuses a topology that cannot run.

Read Chapter 1 before this chapter. Read Working with ndarray and Installation and Feature Flags too, because the neural_network feature gates this whole module. Read Error Handling as well, because layer and loss constructors return Result. A chain model has this end-to-end shape:

use rustyml::neural_network::{
    Shape,
    layers::{Activation, Dense},
    losses::MeanSquaredError,
    optimizers::Adam,
    sequential::SequentialBuilder,
};
use ndarray::Array;

fn main() {
    let x = Array::ones((8, 4)).into_dyn(); // 8 samples, 4 features
    let y = Array::ones((8, 1)).into_dyn(); // 8 samples, 1 target

    // The builder collects the layers, and `build` gives every one of them the shape that
    // reaches it. A stack that does not agree is refused here, before any data moves.
    let mut model = SequentialBuilder::new()
        .add(Dense::new(16, Activation::ReLU).unwrap())
        .add(Dense::new(1, Activation::Linear).unwrap())
        .build(&Shape::known(x.shape()))
        .unwrap();
    model.compile(
        Adam::new(0.001, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    model.fit(&x, &y, 5).unwrap();
    let preds = model.predict(&x).unwrap();
    println!("prediction shape: {:?}", preds.shape());
}

A graph model wires the same layers by node. GraphBuilder::input declares an inlet, add puts 1 layer in the model and calls it, and build names the outlets. Each call takes the nodes that feed it, so a residual connection is 1 extra name in that list. Graph::fit, Graph::predict, and Graph::evaluate take 1 tensor per inlet and give 1 tensor per outlet:

use ndarray::Array;
use rustyml::prelude::*;

fn main() {
    let x = Array::ones((8, 4)).into_dyn();
    let y = Array::ones((8, 1)).into_dyn();

    let mut builder = GraphBuilder::new();
    let input = builder.input(Shape::known(&[8, 4]));
    let hidden = builder.add(Dense::new(4, Activation::ReLU).unwrap(), &[input]);

    // `Add` reads 2 nodes, so the block adds its own input to what the layer computed.
    let residual = builder.add(Add::new(), &[input, hidden]);
    let head = builder.add(Dense::new(1, Activation::Linear).unwrap(), &[residual]);
    let mut model = builder.build(&[head]).unwrap();

    model.compile(
        Adam::new(0.001, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );
    model.fit(&[&x], &[&y], 5).unwrap();

    let preds = model.predict(&[&x]).unwrap();
    println!("prediction shape: {:?}", preds[0].shape());
}

Add above belongs to the merge family, which a graph model needs and a chain model cannot use. The family has 7 layers: Add, Subtract, Multiply, Average, Maximum, Minimum, and Concatenate. The first 6 reduce their inputs element by element, and Concatenate joins them along 1 axis. 3.10 covers all 7, and the shape rule of each.

Average is the 1 name that 2 categories of the crate give to an item. The root of the prelude therefore keeps metrics::Average. The merge layer is prelude::neural_network::Average, or neural_network::layers::Average.

The Sequential Model is the container that turns a stack of layers into a trainable network. SequentialBuilder collects the layers and build allocates them. The built model owns the training loop, the optimizer, and the loss. It exposes compile, fit, train_batch, evaluate, predict, summary, input_shape, output_shape, and weight save and load. It also holds the batch-shuffle seed (set_seed) and the learning-rate pair (learning_rate and set_learning_rate). Read this section first, even for a model that needs 1 specific layer alone.

Dense Layers and Activations covers the fully connected layer, the main layer of tabular models and the output stage of most networks. The layer contracts the last axis of an input of rank 2 or more, so 1 kernel also transforms every timestep of a sequence. It also covers the Activation enum, a 12-variant Copy type such as ReLU, Sigmoid, Tanh, Softmax { axis }, and Linear. An activation folds into a layer, or stands as a layer of its own. Read this section second. Everything after it assumes that a Dense layer declares its units alone, and reads its input width from the shape the build hands it.

Loss Functions is the objective half of compile. It covers mean squared error and mean absolute error for regression, and binary, categorical, and sparse-categorical cross-entropy for classification. Read this section closely.

Its averaging conventions differ on purpose: some average per element, and others average per prediction site. A switch between the 2 conventions rescales the effective learning rate. CategoricalCrossEntropy and SparseCategoricalCrossEntropy take a from_logits flag that decides whether the model needs a Softmax on its output. BinaryCrossEntropy has no such flag, and it always expects a probability in (0, 1).

Optimizers is the update half of compile. It covers SGD with momentum, Adam, AdamW, RMSprop, and AdaGrad. This section also covers clip-by-global-norm (global_clipnorm), coupled and decoupled weight decay, and mid-training learning-rate scheduling. Sections 3.1 through 3.4 together give a complete, trainable feed-forward network.

The remaining sections add specialized layers. All of them plug into the same Sequential, and all of them plug into a Graph as well.

Convolutional Layers provides 1D/2D/3D convolution, the transposed convolution that runs it backwards to grow a tensor, plus depthwise and separable variants for spatial data. All 10 take a dilation_rate that widens the receptive field for free. Conv1D also takes causal padding, for a model that must not look ahead.

Pooling Layers provides parameter-free max and average downsampling, plus their global variants, and closes with the UpSampling1D/2D/3D layers that regrow the spatial axes.

Recurrent Layers covers SimpleRNN, LSTM, and GRU for sequences. The return_sequences and go_backwards flags let them stack and run in reverse. A bidirectional model needs the Reverse layer as well. A go_backwards branch returns its states in processing order, and Reverse puts them back into input order before a merge. That section closes with the Embedding layer that turns word indices into the vectors those layers read.

Regularization and Normalization Layers covers dropout, Gaussian noise, and batch, layer, group, instance, and unit normalization. The dropout family includes spatial dropout and a noise_shape that shares 1 draw across an axis. That section also covers the parameter-free Rescaling layer that conditions an input in front of the first weight.

Only the last of those depends on mode. Such a layer behaves differently in fit than in predict, and the model picks the mode through the context it builds.

Saving and Loading Weights covers persistence. save_to_path and load_from_path persist weights only, in postcard binary format. They do not persist the architecture. Build the identical layer stack in code, then load the weights into it. Read this section once a model is worth keeping. See Model Persistence in Depth for the format details and version caveats.

Graph Models and Merge Layers closes the chapter. It covers GraphBuilder, the 5 calls that assemble a topology, and the refusals that build gives before any data moves. It also covers the 7 merge layers and the shape rule each one applies. Read this section for a model that a chain cannot express.