在移动和嵌入式设备上运行推理具有挑战性，因为资源限制很严格；人们必须在严格的功耗要求下使用有限的硬件。在本文中，我们想展示 TensorFlow Lite (TFLite) 内存使用方面的改进，这些改进使其更适合在边缘进行推理。
通常，神经网络可以被认为是一个计算图，它由运算符（如 CONV_2D 或 FULLY_CONNECTED）和存储中间计算结果的张量（称为**中间张量**）组成。这些中间张量通常会预先分配，以减少推理延迟，但代价是内存空间。但是，如果以幼稚的方式实现，这种代价在资源受限的环境中不可小觑；它会占用大量的空间，有时甚至比模型本身大几倍。例如，MobileNet v2 中的中间张量占用了 26MB 内存（图 1），大约是模型本身的两倍。
 |
| 图 1. MobileNet v2 的中间张量（顶部）及其大小映射到二维内存空间（底部）。如果每个中间张量使用专用内存缓冲区（以 65 种不同的颜色表示），它们将占用约 26MB 的运行时内存。 |
好消息是，由于数据依赖性分析，这些中间张量不需要同时存在于内存中。这使我们能够重用中间张量的内存缓冲区，并减少推理引擎的总内存占用。如果网络的形状是一个简单的链，则两个大的内存缓冲区就足够了，因为它们可以在整个网络中交替使用。但是，对于形成复杂图形的任意网络，这种**NP 完全** **资源分配问题**需要一个好的近似算法。
我们为此问题设计了一系列不同的近似算法，它们根据神经网络和内存缓冲区的特性而表现不同，但它们都使用了一个共同点：**张量使用记录**。中间张量的张量使用记录是一个辅助数据结构，它包含有关张量大小以及它在网络给定执行计划中首次和最后使用时间的的信息。借助这些记录，内存管理器能够在网络执行过程中的任何时刻计算中间张量使用情况，并优化其运行时内存，以尽可能减少占用空间。
在 TFLite GPU OpenGL 后端 中，我们为这些中间张量使用 GL 纹理。这些纹理带有一些有趣的限制：（a）纹理的大小在创建后不能修改，（b）在给定时间，只有一个着色器程序可以独占访问纹理对象。在**共享内存缓冲区对象**模式下，目标是**最小化对象池中所有创建的共享内存缓冲区对象的总大小**。这种优化类似于众所周知的 寄存器分配问题，但由于每个对象的尺寸可变，因此复杂得多。
借助上述张量使用记录，我们设计了 5 种不同的算法，如表 1 所示。除了 Min-Cost Flow 之外，它们都是 贪婪算法，每种算法都使用不同的启发式方法，但仍然能够达到或非常接近 理论下限。某些算法根据网络拓扑结构的表现优于其他算法，但总的来说，GREEDY_BY_SIZE_IMPROVED 和 GREEDY_BY_BREADTH 会生成占用内存最小的对象分配结果。
 |
| 表 1. 共享对象策略的内存占用（单位：MB；最佳结果以绿色突出显示）。前 5 行是我们的策略，后 2 行作为基线（下限表示最佳可能数字的近似值，可能无法实现，而朴素表示每个中间张量分配其自身内存缓冲区时的最差可能数字）。 |
回到我们开头的示例，GREEDY_BY_BREADTH 在 MobileNet v2 上表现最佳，它利用了每个运算符的广度，即运算符配置文件中所有张量的总和。图 2，尤其是与图 1 相比，突出了使用智能内存管理器可以获得多大的收益。
 |
| 图 2. MobileNet v2 的中间张量（顶部）及其大小映射到二维内存空间（底部）。如果中间张量共享内存缓冲区（以 4 种不同的颜色表示），它们将仅占用约 7MB 的运行时内存。 |
对于在 CPU 上运行的 TFLite，适用于 GL 纹理的内存缓冲区属性不适用。因此，通常的做法是在预先分配一个巨大的内存区域，并让所有读取器和写入器共享该区域，它们通过给定的偏移量访问该区域，该偏移量不会与其他读取和写入冲突。这种**内存偏移量计算**方法的目标是**最小化内存区域的大小**。
我们为此优化问题设计了 3 种不同的算法，并研究了之前的工作（Sekiyama 等人 2018 年的带状打包）。与共享对象方法类似，某些算法根据网络的表现优于其他算法，如表 2 所示。从这项研究中得出的一个结论是，**偏移量计算方法总体上比共享对象方法具有更小的占用空间**，因此，如果适用，应优先选择前者而不是后者。
 |
| 表 2. 偏移量计算策略的内存占用（单位：MB；最佳结果以绿色突出显示）。前 3 行是我们的策略，下一行是之前的工作，最后 2 行作为基线（下限表示最佳可能数字的近似值，可能无法实现，而朴素表示每个中间张量分配其自身内存缓冲区时的最差可能数字）。 |
这些针对 CPU 和 GPU 的内存优化已在最近几个稳定的 TFLite 版本中默认提供，并且已被证明对支持更多要求苛刻的最新模型（如 MobileBERT）很有价值。您可以通过查看 GPU 实现 和 CPU 实现 直接找到有关实现的更多详细信息。
Matthias Grundmann、Jared Duke、Sarah Sirajuddin，并特别感谢 Andrei Kulik 的最初头脑风暴以及 Terry Heo 对 TFLite 中最终实现的贡献。
