ShrinkRay

ShrinkRay

ShrinkRay is a powerful tool designed for quantizing and compressing small neural networks, targeting microcontrollers such as the ESP32, STM32, nRF52, RP2040, and ATmega. With a process reminiscent of HandBrake for TinyML, this solution enables seamless conversion of neural network models from frameworks like Keras, TFLite, ONNX, and scikit-learn into chip-ready formats, ensuring efficient deployment on your desired hardware.

Key Features

• Verdict System: Receive a preemptive verdict (✅ FITS, ⚠️ TIGHT, ❌ TOO BIG) indicating whether the quantized model will fit within the target chip's flash and RAM.
• Model Compatibility: Supports conversion from Keras or TensorFlow to TFLite formats (int8/int16/float), generating C byte arrays for various runtime environments including LiteRT-M and CMSIS-NN.
• Scikit-learn Support: Easily generate a single portable C99 header from scikit-learn models using emlearn, compatible even with ATmega328P microcontrollers.
• ONNX Compatibility: Convert ONNX models to the .espdl format utilizing ESP-PPQ as an optional feature.
• Local Operation: Operate 100% locally with no network calls, ensuring privacy and security—there is no telemetry or account requirement.

Usage Example

A quick start to using ShrinkRay is straightforward:

pip install shrinkray-cli

shrinkray chips                  # Access the 12-chip database
shrinkray inspect model.keras      # Review model parameters, size, and performance metrics
shrinkray convert model.keras --target esp32s3 --method int8 --data cal.npy

The output includes details such as:

wrote outputs/model_int8.tflite (2.7 KiB)
wrote outputs/model_int8.cc
wrote outputs/model_int8.h
accuracy max-abs-error 0.002717, top-1 agreement 100.0%
esp32s3: ✅ FITS — flash 2.7 KiB / 8192.0 KiB (0.0%), est. RAM 0.1 KiB / 512.0 KiB (0.0%)
wrote outputs/report.md
wrote outputs/report.html

With the inclusion of calibration data (formatted as a small float32 array), users can enable full-integer quantization along with an accuracy check, enhancing model performance validation options.

Command Overview

Customization Options

Users can create a shrinkray.toml configuration file in their project directory to simplify command usage and avoid repeating flags throughout the CLI process. This flexibility enhances usability and streamlines workflows.

Chip Database

The tool is preloaded with a comprehensive database of 12 chips, each characterized by specifications like flash size, SRAM, and clock speed. This feature allows users to expand compatibility by adding new custom chip definitions according to their specific requirements.

Development and Testing

For developers looking to contribute to the ShrinkRay project, installation of dependencies is facilitated via standard Python package management practices. The tool is extensively tested across various platforms to ensure reliability and compatibility.

In summary, ShrinkRay serves as a vital resource for developers looking to efficiently deploy machine learning models on constrained environments, ensuring optimal performance and fit for microcontroller applications.