256KB 内存下的设备端训练
计算机视觉与模式识别
2024-04-04 v4
摘要
设备端训练通过微调预训练模型，使模型能够适应从传感器收集的新数据。用户可受益于定制化 AI 模型，而无需将数据转移至云端，从而保护隐私。然而，训练内存消耗对于具有微小内存资源的 IoT 设备而言高得令人望而却步。我们提出一种算法-系统协同设计框架，使设备端训练仅需 256KB 内存即可实现。设备端训练面临两个独特挑战：（1）神经网络的量化图因低比特精度和缺乏归一化而难以优化；（2）有限的硬件资源不允许完整的反向传播。为应对优化困难，我们提出量化感知缩放以校准梯度尺度并稳定 8 位量化训练。为减少内存占用，我们提出稀疏更新以跳过较不重要层和子张量的梯度计算。算法创新由一个轻量训练系统 Tiny Training Engine 实现，其剪枝反向计算图以支持稀疏更新，并将运行时自动微分卸载至编译时。我们的框架是首个在 256KB SRAM 和 1MB Flash 下无需辅助内存实现卷积神经网络微小设备端训练的解决方案，使用的内存不到 PyTorch 和 TensorFlow 的 1/1000，同时在 tinyML 应用 VWW 上匹配精度。我们的研究使 IoT 设备不仅能执行推理，还能持续适应新数据以进行设备端终身学习。视频演示可在此处找到：https://youtu.be/0pUFZYdoMY8。
引用
@article{arxiv.2206.15472,
title = {On-Device Training Under 256KB Memory},
author = {Ji Lin and Ligeng Zhu and Wei-Ming Chen and Wei-Chen Wang and Chuang Gan and Song Han},
journal= {arXiv preprint arXiv:2206.15472},
year = {2024}
}
备注
NeurIPS 2022