logo

Deep Learning in Rust: Burn 0.4.0 released and plans for 2023

166 points ksyiros
View on Reddit ↗

Hi everyone!

Last year, I announced Burn (https://github.com/burn-rs/burn), a new deep learning framework written in Rust. The response from the community was very positive, which encouraged me to continue working on the project and solidify the core architecture.

I'm happy to share that the backend refactor has been completed and it's now easier than ever to implement a backend for Burn! Each backend implementation lives in its own crate, making it possible to create third party backends. I have put a lot of effort into making the process as easy as possible, and even the tests are generated to ensure that all operations are correctly implemented.

This is the basic building block that allows creating a backend decorator that can optimize the computational graph dynamically. This could potentially make Burn the fastest deep learning framework with dynamic graphs that are comparable to frameworks with statically compiled graphs, while still maintaining the flexibility and improved developer experience that dynamic graphs offer.

A lot of rough edges were also improved, including but not limited to:

  • Refactoring how metrics work to allow any metric to be used by adapting the model outputs and leveraging the Rust type system to ensure type correctness.
  • Decoupling the optimizer from the module trait by introducing a ModuleVisitor. This new trait allows any struct to visit all tensor parameters in a module, which the optimizer can use to execute its step.
  • Cleaning up how states are saved, removing a lot of code and using the std io API to compress and save states with reduced memory usage.

Also, a lot of new features have been implemented, such as the ADAM optimizer, transformer modules, and more. The next focus will be on adding vision-related modules like convolution and pooling layers. If you are interested in helping out or just want to follow the development more closely, I'm creating a Discord channel at https://discord.gg/KcVGzmCcWj. I'm confident that 2023 will be a good year for Burn and I hope that Rust will start becoming a strong contender in the realm of scientific computing. Happy new year everyone!