TFLite micro内存管理与分配策略的优化

TFLite micro(TFLm)是当前在微控制器平台上流行的神经网络推理框架。本文分析了TFLm在推理模型时的内存管理机制与分配策略,以及其在使用内存时的局限性。当前TFLm仅支持使用单块内存(Tensor Arena)来保存模型推理所需的中间结果,本文扩展TFLm的内存管理以支持使用多块不连续且访问性能有巨大差异的内存,还给可以重叠的tensor分配相同的内存。通过这样的改进,既把数据流量更多地引到片上快速内存中,又降低了峰值内存的占用。通过在i.MX RT1170上的实验数据表明,本文策略对于含有快速片上RAM(以DTCM为代表)的微控制器,能大大提高片上快速RAM的利用率,显著缓解存储器带宽带来的瓶颈,使推理时间缩短至一半以上。