Skip to content

DESIGN.md

Latest commit

 

History

History
69 lines (55 loc) · 4.26 KB

DESIGN.md

File metadata and controls

69 lines (55 loc) · 4.26 KB
Raw

While iterating and refactoring this project, I had the following in mind.

Overview of Goals and Philosophy

  • Strong Type Safety: Through Rust, leverage type-driven design features and patterns (const generics, zero-cost abstractions, typestate, etc) to catch mismatches and human errors at compile time; higly reflective of "parse, don't validate".
  • Ergonomic API: Provide a clear and concise API for model construction, training, and testing.
  • Constrained Modularity: Allow legal mixing and matching of layers without committing to a rigid static pipeline.
  • Performance: Use zero-cost abstractions, heap-allocated (Box-ed) arrays, and row-major order data layout.
  • Minimal Dependencies: Keep external dependencies to a minimum, relying on crates and the standard library at very select times; e.g., the rand crate.

These were the main design philosophies I kept in mind through the project. For more detail, consult the remaining of this document.

Linear Algebra

Since I wanted to make this project absolutely from the ground up, I would end up making some quick linear algebra utilities.

  1. Store vectors and matrices in row-major order for cache-friendly access during matrix-vector multiplication.
  2. Represent fixed-size arrays behind a Box<[T; N]> or Box<[[T; N]; M]> to minimize stack usage while retaining compile-time size checks.
pub struct Matrix<T, const N: usize, const M: usize> {
  entries: Box<[[T; N]; M]>,
}
  1. Generate random vector and matrix using the rand crate.

API design

There was a lot of thought put into API design. For example, consider where weights associated between two layers would be stored. Although it might not have seemed logical to associate weights with one layer, this ended up being the case to avoid keeping track of more indices, slightly improving performance and code readability.

Initialization

let network = ModelBuilder::new()
  .input(dim!(128))
  .hidden(dense!(64).activation(ReLU))
  .output(dim!(1));

This would be the dream to walk towards. Firstly, the macro might be yelling at you. Let me explain. Since I wanted to put type-safety as the highest priority, I would need to have called an initializer function like LayerBuilder::dense::<128, 128>() no matter what because of const generics, which just looks ugly. So, a macro needed to be invoked in this context, which I call dense!(n).

However, the API ended up looking like this:

let network = ModelBuilder::new()
  .input(dim!(128))
  .hidden(dense!(64, 64).activation(ReLU))
  .output(dim!(1));

The stable Rust compiler can unfortunately not do inferences with const generics yet, so you have to be very explicit with the input and output dimensions of each layer. There have been numerous discussions on const generic inferences before, and it seems to be pretty close to becoming stabilized; this is something I would go back and fix once the Rust team stabilizes it.

For now, to use the cleaner syntax, you must enable the unstable library feature and write #![feature(generic_arg_infer)] on Nightly builds.

Even in Nightly, I'm thinking of making the API look like this:

let network = ModelBuilder::new()
  .input(dim!(128))
  .hidden(LayerBuilder::dense(dim!(64)).activation(ReLU))
  .output(dim!(1));

Although it's significantly more terse, I feel it makes the structure of what you're building so much more obvious.

Training/Testing

let training_data = [(1.0, 2.0), /* more */].map(DataSample::from);
network.fit(&training_data, TrainConfig::default());

let testing_data = [(3.0, 15.0), /* more */].map(DataSample::from);
let y = network.predict(&[3.0]);
dbg!(y);

Note

This part is still being refactored.