BitNetMCU: 为低端微控制器打造的高精度低比特量化神经网络

BitNetMCU 是一个专注于低比特量化神经网络训练和推理的开源项目,旨在为 CH32V003 等低端微控制器提供高效的神经网络解决方案。通过量化感知训练(QAT)和精心调优的模型结构与推理代码,该项目在 16x16 MNIST 数据集上实现了超过 99% 的测试精度,同时仅使用 2KB RAM 和 16KB Flash,无需任何乘法指令。这一成果充分展示了在资源受限的微控制器上部署高精度神经网络的可能性。

项目概述

BitNetMCU 项目的核心目标是探索在低端微控制器上实现高性能神经网络的可行性。项目采用 PyTorch 构建训练流程,确保了良好的可移植性;同时,推理引擎使用 ANSI C 实现,可轻松移植到各种微控制器平台。

项目特性

低比特量化: 支持二值、三值、2比特、4比特和8比特等多种量化方案,适应不同的精度和资源需求。
量化感知训练: 采用量化感知训练技术,在训练过程中模拟量化效果,提高模型在低比特表示下的性能。
无乘法推理: 针对无硬件乘法器的微控制器优化,实现高效的推理过程。
灵活的训练流程: 提供可配置的训练参数,支持模型结构和训练策略的快速调整。
多平台支持: 训练部分基于 PyTorch 实现,推理引擎使用 ANSI C 编写,确保了良好的跨平台兼容性。

项目结构

BitNetMCU 项目采用模块化设计,主要包含以下组件:

BitNetMCU/
│
├── docs/                      # 项目文档
├── mcu/                       # CH32V003 微控制器相关代码
├── modeldata/                 # 预训练模型存储
│
├── BitNetMCU.py               # PyTorch 模型和 QAT 类定义
├── BitNetMCU_inference.c      # C 语言推理代码 
├── BitNetMCU_inference.h      # C 语言推理代码头文件
├── BitNetMCU_MNIST_test.c     # MNIST 数据集测试脚本
├── BitNetMCU_MNIST_test_data.h# MNIST 测试数据头文件(生成)
├── BitNetMCU_model.h          # C 语言模型数据头文件(生成)
├── exportquant.py             # 模型量化导出脚本
├── test_inference.py          # C 语言推理实现测试脚本
├── training.py                # 神经网络训练脚本
└── trainingparameters.yaml    # 训练参数配置文件

这种结构设计使得项目的各个组件职责清晰,便于开发者理解和扩展。

训练流程

BitNetMCU 项目的训练流程设计灵活,可以满足不同的实验需求:

配置: 通过修改 trainingparameters.yaml 文件设置训练超参数。
模型训练: 使用 training.py 脚本训练模型,生成的权重以 .pth 格式保存在 modeldata/ 目录下。
模型量化: 运行 exportquant.py 脚本将模型转换为量化格式,并导出为 C 语言头文件 BitNetMCU_model.h。
C 模型测试: 编译并执行 BitNetMCU_MNIST_test.c 进行推理测试。
Python vs C 模型对比: 使用 test-inference.py 脚本对比 C 语言和 Python 模型在完整 MNIST 测试集上的性能。
微控制器测试: 按照 mcu/readme.md 中的说明在实际硬件上进行测试。

这一流程设计既方便了模型的快速迭代和优化,又确保了最终部署到微控制器上的模型性能。

技术亮点

多种量化方案:
- 二值化(1比特)
- 三值化(1.58比特)
- 2比特对称量化
- 4比特对称和非对称量化
- FP1.3.0 量化(支持无乘法推理)
- NormalFloat4 (NF4) 量化
OCTAV 最优裁剪算法: 用于计算最优裁剪和量化参数,提高量化模型的精度。
灵活的权重缩放: 支持 PerTensor 和 PerOutput 两种权重缩放方式,适应不同的硬件特性。
数据增强: 支持旋转和仿射变换等数据增强技术,提高模型泛化能力。
动态学习率调整: 实现了 StepLR 和 CosineAnnealingLR 两种学习率调度策略。

项目进展

BitNetMCU 项目持续更新,不断引入新的特性和优化:

2024年4月24日: 首次发布,支持二值、三值、2比特、4比特和8比特量化。
2024年5月2日: 发布 0.1a 版本。
2024年5月8日: 添加 FP1.3.0 量化,实现 98.9% 精度的无乘法推理。
2024年5月11日: 修复 Linux 兼容性问题。
2024年5月19日: 增加非对称4比特量化支持。
2024年5月20日: 引入 quantscale 超参数,优化权重缩放。
2024年5月26日: 发布 0.2a 版本。
2024年7月19日: 引入 OCTAV 算法,优化裁剪和量化参数。
2024年7月26日: 增加 NormalFloat4 (NF4) 量化支持。

应用前景

BitNetMCU 项目为低端微控制器上的神经网络应用开辟了新的可能性。以下是一些潜在的应用场景:

嵌入式视觉: 在资源受限的设备上实现简单的图像分类和识别任务。
智能传感器: 提升传感器数据处理能力,实现本地智能决策。
预测性维护: 在工业设备中集成简单的异常检测模型。
语音关键词识别: 实现低功耗的语音唤醒功能。
手势识别: 在可穿戴设备中实现简单的手势交互。

开发者社区

BitNetMCU 是一个开源项目,欢迎开发者参与贡献。项目在 GitHub 上获得了 221 颗星和 20 次分叉,反映了社区对该项目的兴趣。开发者可以通过以下方式参与:

提交 Issue 报告 bug 或提出新功能建议
提交 Pull Request 贡献代码
在 Discussions 中讨论技术问题和应用思路
为项目文档做出贡献,帮助新用户更好地理解和使用 BitNetMCU

未来展望

BitNetMCU 项目虽然已经取得了显著成果,但仍有很大的发展空间:

扩展模型支持: 除 MNIST 外,探索在其他数据集上的应用。
优化量化算法: 继续研究更高效的量化方法,进一步提升精度和效率的平衡。
硬件加速: 探索利用微控制器特定硬件特性进行推理加速的可能性。
端到端优化: 结合硬件特性,从模型设计到推理实现的全流程优化。
跨平台支持: 扩展到更多种类的微控制器平台。

BitNetMCU 项目展示了在资源受限环境下实现高性能神经网络的可能性,为嵌入式 AI 和边缘计算领域提供了宝贵的参考。随着项目的不断发展和社区的积极参与,我们有理由期待 BitNetMCU 在推动低端设备 AI 应用方面发挥更大的作用。

结语

BitNetMCU 项目通过创新的量化技术和优化策略,成功地在低端微控制器上实现了高精度的神经网络推理。这一成果不仅展示了嵌入式 AI 的巨大潜力,也为解决资源受限设备上的复杂问题提供了新的思路。随着物联网和边缘计算的快速发展,像 BitNetMCU 这样的项目将在推动技术创新和应用落地方面发挥越来越重要的作用。我们期待看到更多开发者和研究者加入到这一激动人心的领域,共同探索人工智能在微控制器世界中的无限可能。