Zffnn - Comptime Neural Network Inference Engine

Showcase
, ,
1

I wanted to dig into what comptime is capable of, so I got the idea to try and inspect static neural networks and their weights in Zig. This developed into me creating my first sort of library, with a user oriented API and all, and ended with super fast, static inference binaries for (mostly) arbitrary neural networks. That being said, forgive me for any shortcomings, I’m still learning so any comments or concerns are welcome.

The library is built upon 3 fundamental, compile-time defined types that possess runtime functions:

Matrix, Layer, and NN; with each type being built upon the other.

The network shape is defined by the user like:

const definition: []const struct { usize, Activation } = &.{
    .{784, .none},     // input layer (no activation)
    .{128, .relu},     // hidden
    .{10, .softmax},   // output
};
const Net = NN(definition, batch_size);

This is used to create the network and its internals without needing heap or runtime allocations.

In collaboration with Zig’s aggressive optimization in ReleaseFast, the entire scope of the program determined at compile time allows for the library to produce extremely performant binaries for models in it’s sweet spot (small models at small batch size that don’t pressure memory too heavily).

A few things of note/greater questions:

  • I was able to configure comptime paramater embedding through some sustained effort; this was an integral part of the practicality of this project and proved to be one of the more difficult parts to get right. If you have any ideas for how to do this idiomatically I’d appreciate them.

  • Zig’s @Vector types were a big help in effectiveness, but the changes to access in 0.16.0 has forced me to make some changes that resulted in decreases in performance benchmarks across the board.

Sorry, long release post, but I’ve been coding in Zig for about a year and I’m new to this community, so if you took the time to read this thanks.
If you wanna check out the project:
https://github.com/krypticlogan/zffnn

4 Likes