BitNetMCU Slims Down TinyML Models

The major releases of new generative artificial intelligence (AI) tools, like large language models and text-to-image generators, tend to grab all of the headlines these days. But while the organizations behind these systems are still trying to work out exactly what practical uses they have, and how they can make money with them, smaller-scale AI tools are already transforming a number of industries. TinyML is already being used to keep manufacturing equipment in tip-top shape, monitor the environment, and keep watch over people’s health.

New tinyML applications are being enabled all the time by advances in hardware that serve to shrink its size, cost, and energy consumption. Furthermore, algorithmic advancements are allowing more powerful models to run with less and less resources. Most models are trained primarily with accuracy in mind, so many of these strategies revolve around reducing the size of the model, in one way or another, without any significant drop in accuracy.

A new framework called BitNetMCU seeks to simplify the process of training and shrinking neural networks such that they can run on even the lowest-end of microcontrollers. The initial focus of the project is to support low-end RISC-V microcontrollers like the CH32V003, which only has 2 kilobytes of RAM and 16 kilobytes of flash memory. Moreover, this chip lacks a multiply in its instruction set (which is kind of important for the matrix multiplications that are a key component of neural network computations). However, the CH32V003 only costs about 15 cents in small quantities, so there is good reason to target such a chip — especially for large, distributed networks of sensors.

The BitNetMCU pipeline consists of a set of Python scripts that use PyTorch on the backend to train a neural network, then perform a low-bit quantization to slash it down to a size that can fit on a tiny microcontroller. After training, a set of other scripts are available for testing the model’s performance. And once everything is looking good, there is another utility available to assist in deploying the model to a CH32V003 microcontroller — although, in principle, it should not be too difficult to deploy to other microcontrollers as well. The inference engine was implemented in ANSI C, so support will be available on a wide range of hardware platforms.

To test the system, an experiment was conducted to see how well a handwritten digit recognition model would perform against the 16x16 MNIST dataset. Because the model was heavily quantized before deployment, it was able to fit within the tiny amount of RAM and flash memory that are available to the CH32V003. As for the chip’s other big shortcoming — the lack of a multiply instruction — BitNetMCU makes up for it through the use of multiple additions. Testing of the model showed that it had achieved an accuracy level of better than 99 percent.

Source code for the project is available on GitHub, and there is some fairly good documentation as well. So, if you have a few dimes burning a hole in your pocket, BitNetMCU could be a fun way to get started in the world of tinyML.