Compile-time creation of neural networks with Rust
This is for now just a showcase project of what can be done with const generics
introduced in Rust 1.51. There is not a single usage of vec
in this project (as of today).
GAMMA allows the developer to build neural networks at compile-time, with preallocated arrays with well defined sizes. Aside from the performance improvement at runtime, another important benefit is that any possible mistake with the layout of the neural network, for example mismatching the inputs/outputs in the chain of layers, will be raised at compilation time.
This magic is accomplished thanks to two awesome Rust features:
const generics
: The layer weights are defined as multi-dimensional arrays with generic sizes. Before this feature was introduced the only option was to usevec
or go crazy and define different layer types for each possible number of weights!derive macro
: It is impossible to define an array or any other iterable of layers because it is an hetereogeneous set (different number of weights for each layer). To perform the forward pass you need to chain all the layers and propagate the input up to the lastest layer. TheNeuralNetwork
derive macro defines theforward
method at compile-time, doing exactly that without any iteration.
Add this to your Cargo.toml
:
[dependencies]
gamma = "0.1"
gamma_derive = "0.1"
And this is a very simple example to get you started:
use rand::distributions::Uniform;
use gamma::{activations::relu, layers::Dense, NeuralNetwork};
use gamma_derive::NeuralNetwork;
// Builds a neural network with 2 inputs and 1 output
// Made of 3 feed forward layers, you can have as many as you want and with any name
#[derive(NeuralNetwork, Debug)]
struct MyNetwork {
input: Dense<2, 4>, // <# inputs, # outputs>
hidden: Dense<4, 2>,
output: Dense<2, 1>,
}
impl MyNetwork {
// Initialize layer weights with a uniform distribution and set ReLU as activation function
fn new() -> Self {
let mut rng = rand::thread_rng();
let dist = Uniform::from(-1.0..=1.0);
MyNetwork {
input: Dense::random(&mut rng, &dist, relu),
hidden: Dense::random(&mut rng, &dist, relu),
output: Dense::random(&mut rng, &dist, relu),
}
}
}
fn main() {
// Init the weights and perform a forward pass
let nn = MyNetwork::new();
println!("{:#?}", nn);
let input = [0.0, 1.0];
println!("Input: {:#?}", input);
let output = nn.forward(input);
println!("Output: {:#?}", output);
}
You may wonder how the forward
method works. The NeuralNetwork
derive macro defines it for you, and it looks like this for this particular example:
fn forward(&self, input: [f32; 2]) -> [f32; 1] {
self.output.forward(self.hidden.forward(self.input.forward[input]))
}
Note how the forward method expects two input values because that's what the first (input
) layer expects, and returns one single value because that's what the last layer (output
) returns.
- Compile-time neural network consistency check
- Docs, CI/CD & Benchmarks
- Backward pass
- More layer types (convolution, dropout, lstm...)
- More activation functions (sigmoid, softmax...)
- Maaaybeee, CPU and/or GPU concurrency
If you find a vulnerability, bug or would like a new feature, open a new issue.
To introduce your changes into the codebase, submit a Pull Request.
Many thanks!
GAMMA is distributed under the terms of both the MIT license and the Apache License (Version 2.0).
See LICENSE-APACHE and LICENSE-MIT, and COPYRIGHT for details.