Building a Deep Learning Framework in Rust

Creating a deep learning framework from scratch is a monumental task, but one developer took on the challenge and built a graph-based framework in pure Rust. In the first two parts of this series, the developer argued for the importance of such a project and wrote the GPU backend on wgpu. However, the code was not yet ready for other humans to use.

Phase 1: Cleanup and Consistency

The first step in making the code more usable was to clean up the existing codebase. The developer ran cargo build --release --all-targets and was greeted with 142 warnings. After some investigation, it was discovered that the binary was recompiling the entire library as a separate crate, resulting in many false-positive warnings. By replacing the mod declarations with use rustyasg::*, the developer was able to eliminate around 100 false-positive warnings. The rest of the warnings were addressed through actual cleanup, including deprecating rand::thread_rng and removing unused imports.

Phase 2: Declarative Layer API

The next step was to create a declarative layer API. The existing code used string-matching to figure out parameter shapes, which was not only cumbersome but also prone to errors. The developer decided to make layers own their shape information, with every nn::* constructor taking dimensions and self-registering with GraphContext. This change required a significant overhaul of the existing code, but the result was a much more user-friendly API.

Phase 3: GPU Completeness

With the API changes in place, the developer turned their attention to achieving GPU completeness. The existing GPU backend was incomplete, with some operations not supported on the GPU. The developer wrote the necessary WGSL shaders to support these operations, including four shaders for LayerNorm alone.

As the developer looks to the future, they are focused on continuing to improve and expand the framework. With the release of v0.3.1, the framework is now more usable and efficient than ever before. The developer is excited to see how the community will use and contribute to the framework, and is committed to continuing to push the boundaries of what is possible with deep learning in Rust.