256KB内存下的设备端训练

[TinyML项目官网] | [论文] | [网站] | [YouTube] | [B站]

最新消息

如果您希望获取最新动态，请在此 注册 以接收通知。

• [2024年3月6日] 我们发布了关于256KB内存下的设备端训练的新演示视频。
• [2023年10月13日] IEEE CAS杂志刊登了微型机器学习：进展与未来 [专题]。
• [2022年11月28日] 我们的海报展示将于11月30日星期三上午11:00至下午1:00（新奥尔良时间）在J厅702号展位进行。欢迎感兴趣的朋友前来参观！
• [2022年10月4日] 我们关于微型设备端训练的论文被MIT主页重点推荐！
• [2022年9月16日] 我们的论文已被NeurIPS 2022接收！
• [2022年6月30日] 我们的论文已在arXiv上发布。

概述

过去，深度神经网络的训练通常在云端进行。那么，我们能否在边缘设备上进行学习呢？最大的挑战在于内存占用过大。由于内存预算非常有限（比GPU小5万倍），即使仅进行推理也十分困难，更不用说训练了。

在本工作中，我们通过系统-算法协同设计，实现了在256KB内存限制下的设备端训练，所用内存不到PyTorch的千分之一，同时在视觉唤醒词应用上的准确率与之相当。

我们的工作包含三个关键部分，以实现高效的设备端训练：

1. 量化感知缩放与稀疏更新（算法）
2. 编译时自动微分及对稀疏更新的系统支持（系统）
3. 使用TinyEngine后端的代码生成（系统）

如果您对系统的全栈优化感兴趣，可以逐一了解每个步骤；若您只想快速搭建一个小型训练演示，则可以直接参考Tiny Training Engine中的示例设置。

1. 量化感知缩放（QAS）与稀疏更新

为了优化真实的量化图（见下文对比），我们提出了量化感知缩放（QAS）方法，用于自动调整梯度大小，从而有效稳定训练过程，并达到与FP32精度相当的效果。

此外，我们还设计了稀疏层和稀疏张量更新机制，跳过不重要层和子张量的梯度计算。针对不同内存预算，我们基于贡献度分析，自动确定最佳的稀疏化程度。

我们在算法文件夹中提供了评估QAS和稀疏更新准确性的脚本，并将使用预训练模型和稀疏更新方案来进行下一步的编译工作。

2. 编译时自动微分及对稀疏更新的系统支持

我们工作的另一大亮点是Tiny Training Engine（TTE），它将自动微分从运行时移至编译时，并通过代码生成来最小化运行时开销。同时，TTE还支持图剪枝和重排序，以实现稀疏更新，从而将理论上的节省转化为实际的内存占用减少和速度提升。

与编译、自动微分及稀疏更新系统支持相关的代码位于编译文件夹中。这些代码会将PyTorch模型转换为中间表示（IR），在编译时完成自动微分，并应用稀疏更新规则以降低内存使用。最后，经过剪枝的训练图会被转换为JSON文件，便于在MCU上部署。

3. 使用TinyEngine后端的代码生成

获得训练图的JSON表示后，下一步就是按照Tiny Training Engine教程将模型部署到MCU上，从而完成开头展示的演示。我们优化后的内核和协同设计不仅使训练能够在256KB内存限制下进行，还比传统的实现方式（如TF-Lite）更快。

引用

 @inproceedings{lin2022ondevice,
    title     = {On-Device Training Under 256KB Memory},
    author    = {Lin, Ji and Zhu, Ligeng and Chen, Wei-Ming and Wang, Wei-Chen and Gan, Chuang and Han, Song},
    booktitle = {Annual Conference on Neural Information Processing Systems (NeurIPS)},
    year      = {2022}
} 

相关工作

• MCUNet：物联网设备上的微型深度学习（NeurIPS'20）
• MCUNetV2：面向微型深度学习的高效内存占用补丁式推理（NeurIPS'21）
• TinyTL：减少激活值而非可训练参数以实现高效设备端学习（NeurIPS'20）
• 用于微型深度学习的网络增强技术（ICLR'22）

许可证

本仓库采用MIT许可证开源。更多详情请参阅LICENSE文件。

tiny-training 快速上手指南

tiny-training 是一个旨在在极低内存（<256KB）的微控制器（MCU）上进行深度学习模型训练的开源自适应工具。本项目通过“系统 - 算法协同设计”，实现了比 PyTorch 少用 1000 倍内存的端侧训练能力。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04/20.04) 或 macOS。Windows 用户建议使用 WSL2。
• Python 版本: Python 3.7 - 3.9
• 核心依赖:
  • PyTorch (用于模型定义与算法验证)
  • GCC / Clang (用于编译后端)
• 硬件目标 (可选): 如需在物理设备上运行演示，需准备支持 CMSIS-NN 的 MCU 开发板（如 Arduino Nano BLE Sense, STM32 等）及对应的烧录工具链。

提示：国内开发者安装 PyTorch 时，推荐使用清华或中科大镜像源以加速下载。

安装步骤

1. 克隆仓库

   git clone https://github.com/mit-han-lab/tiny-training.git
   cd tiny-training
   

2. 创建虚拟环境并安装依赖

   python3 -m venv venv
   source venv/bin/activate
   
   # 使用国内镜像源安装 PyTorch (CPU 版本示例，如需 GPU 请调整)
   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
   
   # 安装项目其他依赖
   pip install -r requirements.txt
   

3. 验证算法模块 (可选) 如果您仅关注算法验证（量化感知缩放 QAS 和稀疏更新），可直接在 algorithm/ 目录下运行评估脚本：

   cd algorithm
   python eval_qas_sparse.py
   

基本使用

本项目分为算法验证、编译优化和端侧部署三个阶段。对于大多数希望快速构建演示的用户，核心流程是将 PyTorch 模型转换为可在 MCU 上运行的训练图。

1. 模型编译与自动微分 (Compile-time Autodiff)

将预训练的 PyTorch 模型转换为中间表示 (IR)，并在编译时执行自动微分和图剪枝，以生成优化的训练图。

cd compilation

# 示例：将模型转换为 IR 并应用稀疏更新规则
# 请根据实际模型路径调整参数
python compile_training_graph.py \
    --model_path ../algorithm/pretrained_model.pth \
    --output_path training_graph.json \
    --sparsity_ratio 0.5

执行成功后，将在当前目录生成 training_graph.json，该文件包含了优化后的训练计算图。

2. 端侧部署 (Codegen with TinyEngine)

利用生成的 JSON 文件，通过 TinyEngine 后端生成 C 代码并部署到 MCU。详细的部署教程位于独立的 tinyengine 子模块中。

# 进入 TinyEngine 训练教程目录
# 注意：这需要初始化子模块或单独克隆 tinyengine 仓库
git submodule update --init --recursive
cd ../tinyengine/tutorial/training

# 按照以下步骤生成代码并编译固件
# 1. 导入上一步生成的 JSON 文件
# 2. 生成针对特定 MCU 的 C 代码
# 3. 使用 make 编译固件
make BOARD=arduino_nano_ble33_sense

3. 运行演示

将编译好的固件烧录至开发板。设备上电后，即可在 <256KB 内存限制下，利用本地采集的数据（如图像传感器数据）进行实时模型微调训练。

• 观察指标：通过串口监视器查看训练损失下降情况及内存峰值占用。
• 参考视频：YouTube 演示 | Bilibili 演示

---

更多高级配置（如自定义稀疏策略、算子融合优化）请参考仓库中 algorithm/ 和 compilation/ 目录下的详细文档。