TensorFlow Lite源码解析之二（内存管理）

相关阅读
TensorFlow Lite源码解析之一

1. 前言

爱迪生说过，人工智能就是是百分之九十九的数据加上百分之一的算法。毕竟目前人工智能还没有达到T800这种以毁灭人类为己任的终结者级别，归根到底还是一个程序。这么一想，是不是觉得市面上说的AI要统治人类了根本就是危言耸听，对于弱人工智能，我治不住你，难道我的360强力卸载还治不住你？

言归正传。很显然，想要了解一个程序，理解它是怎么管理用于存储数据的内存是一个绕不开的话题。想要了解TensorFlow Lite是如何工作的，我们首先要弄清楚它的Tensors都会Flow去哪。如果Tensor是水的话，那么内存分配的过程就是挖沟的过程。对于TFLite这种用于端测推理的推理引擎，在内存使用上也不能像服务器那么豪横，总之一句话，就是要努力做到又要马儿跑，又不让马吃草。

2. 太长；不看

你的时间非常值钱，打开手机肯定不是为了听我在这哔哔哔，内存分配其实是个繁琐的过程，为了节省时间先说说结论。当然，不着急的话看完也是极好的。

结论就是，TFLite首先会根据每个张量的大小(size)，为它们分配一个偏移地址(offset)，并且保证不会有任何一个张量的数据在错误的时间覆盖任何其他有用的张量。并且，TFlite能够做到用于中间结果的内存总大小不超过最大张量所用空间 * （1+最大分值数），这是最坏的情况，实际情况所使用的空间可能更小。对于分支少的模型，节省的内存非常可观。这极大的缓解的了移动端设备内存的压力。简单地说就是内存复用。

之后，通过最后一个张量的偏移 + 大小 + 必要的首尾填充，得到一个总的内存大小。一次性向系统请求计算出所需的内存，得到一个实际的内存的起始地址。

最后，依次将每个张量的偏移地址加上实际申请得到的内存起始地址，就得到了每个张量数据实际的起始地址，结合张量的大小，最终就可以确定内存中哪块区域属于哪个特定的张量。

3. 详细过程

首先，让我们站在高处，先做个总体的认识。如图1所示，为TFLite进行内存分配的时候的主要的函数调用流程。主要的分配逻辑都由一个名为ArenaPlanner的类负责。

通常，在代码中，如果我们看到有Arena这个词一出现，那么很可能的情况就是程序会通过库文件或者系统调用的方式，直接向内存索取一大片内存，然后再自己分配。在这里，TFLite遵循了这一不成文的规定，在后面的介绍中我们会知道在TFLite中确实是这么做的。

直接向操作系统索取一大片内存再自己做分配的好处显而易见：避免了多次系统调用的开销，因为系统调用确实不便宜，特别在推理引擎这种对时间要求很高的程序中；另外一个就是程序可以根据自己的需要定制分配策略以使得效率更高。这就相当于申请预算，算个大概之后就一次性把下一年度的预算都申请下来，不然在频繁打报告上花费的时间暂且不论，能不能申请下来还是未知数。

在图1中可以看到，TFLite在内存分配上大致可以分成以下几个步骤：

1. 计划阶段。确定整个模型所有的张量中，哪些需要进行内存分配；
2. 计算阶段。计算这些张量所需要的内存大小之和，并确定张量之间的相对地址。就类似于虽然钱还没到手但是我们已经算好了怎么分；
3. 实际分配阶段。一次性向内存申请指定大小的内存，并且将需要的数据进行拷贝。

实际上，TFLite申请的内存分成两块，一块用于储存临时张量以及其他一些临时性数据，另一块则用于存储永久性数据。两块内存不同的是数据的生命周期不同，其他的完全一样。

接下来我们会一一探究每一个过程的的细节。

3.1. 计划阶段

ArenaPlanner通过多个列表来做记录，以辅助内存分配。在计划阶段，也就是在PlanAllocations()调用中，会用到其中三个，他们分别是alloc_node_、dealloc_node_以及refcounts，它们都是std::vector<int>类型的列表，且长度等于当前所在子图所包含的所有张量的数量（length），由于通常整个模型中只有一个子图，因此也可以说这些列表的长度等于整个模型中包含的张量个数。由于每个张量在解析的时候都被赋予了在从0~length-1中唯一的数最为索引，因此这些列表中的每一个元素和模型中的张量是一一对应的。

这三者中前两个是ArenaPlanner的成员变量，后续的操作中需要用到它们在此步骤设置好的值。最后一个refcount是局部变量，主要用于确定dealloc_node_中的元素该如何赋值。

其执行过程如下：

1. 在开始的时候，首先对这三个列表进行初始化：alloc_node_、dealloc_node_的元素都被初始化为0xFFFFFFFF，refcounts的元素都被初始化为0；
2. 为所有类型不是kTfLiteOptionalTensor的张量都都设置引用计数，也就是设置refcounts中与每一个张量对应的元素的值都至少是1；
3. 将alloc_node_中代表输入张量、输出张量以及属于变量张量的元素的值都设置为0；将属于模型中节点的输出的张量的值设置为对应的节点的编号；
4. 如果不需要保留中间结果，则将模型中节点的输入张量在dealloc_node_中对应的元素的值设置成该节点的编号；
5. 对于模型中各个节点所拥有的临时张量，同时将他们在alloc_node_以及dealloc_node_中对应的元素的值设置成该节点的编号。

对于同一个张量（除了模型的输、输出张量以及节点自身的临时张量），它既是前一个张量的输出，同时也是后一个（也可能是多个）节点的输入。因此，对于alloc_node_以及dealloc_node_中相同位置的元素而言，虽然它们对应的是同一个张量，但是值往往不同，其关系至少满足alloc_node_[tensor_index] + 1 = dealloc_node_[tensor_index]，如图2所示。请记住这一点，因为后续会用到这一关系。

这一步完成之后，alloc_node_以及dealloc_node_中同一个张量所对应的元素的值有三种关系：

1. 对于输入、输出以及变量张量：alloc_node_对应的值为0，dealloc_node_对应的值为为0xFFFFFFFF；
2. 对于结算过程中的中间结果张量：如果需要保留中间结果，则alloc_node_对应的值为输出该张量的节点编号，dealloc_node_对应的值为为0xFFFFFFFF；
3. 对于结算过程中的中间结果张量：如果不需要保留中间结果，则alloc_node_对应的值为输出该张量的节点编号，dealloc_node_对应的值为满足alloc_node_[tensor_index] + 1 = dealloc_node_[tensor_index]；划重点，这是TFLite实现内存内存可以使用所用内存空间远小于所有张量内存大小之和的关键所在

<