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实现内存内存可以使用所用内存空间远小于所有张量内存大小之和的关键所在

3.2. 计算阶段

每一个张量所拥有的buffer的信息，通过类ArenaAllocWithUsageInterval来保存，ArenaAllocWithUsageInterval包含偏移（offset）、大小、所属的张量、张量的输入节点、输出节点等信息，如下面代码所示。

struct ArenaAllocWithUsageInterval {
  ArenaAllocWithUsageInterval() { reset(); }

  size_t offset;
  size_t size;
  int32_t tensor;
  int32_t first_node;
  int32_t last_node;

  inline void reset() {
    offset = 0;
    size = 0;
    tensor = -1;
    first_node = -1;
    last_node = -1;
  }

  inline bool operator<(const ArenaAllocWithUsageInterval& other) const {
    return offset < other.offset;
  }
};

在计算阶段要做的就是确定每一个ArenaAllocWithUsageInterval实例中这些属性的值以及总的内存大小。最终TFLite申请内存的时候是按照总的内存大小来申请一块连续的内存，便可以得到这块内存的起始地址，起始地址加上偏移地址就能得到每一个张量所拥有的buffer的起始地址，结合size属性就能确定整个buffer的边界，如下图3所示。

最终，通过计算，包含所有的张量的buffer信息ArenaAllocWithUsageInterval会按照一定的规则有序的存储在一个名叫ordered_allocs_的列表中，另外，也会按照索引号从小到大的顺序保存在allocs_列表中。ordered_allocs_以及allocs_存储的都是ArenaAllocWithUsageInterval的指针。

那么，TFLite具体是怎么确定每个张量的偏移地址的呢？其过程具体如下：

1. 开始的时候，ordered_allocs_会被清空；

2. 为张量的索引值排序，排序的原则为，对于任意两个张量T1以及T2：
   i. 如果T1、T2两个都属于输入张量或者输出张量，那么这T1、T2张量维持他们的顺序不变；
   ii. 如果T1属于输入张量或者输出张量而T2不是，则T1排在T2之前；
   iii. 如果T1、T2都不输入输入张量或者输出张量，则谁需要的内存更多谁排前面；对于需要内存一样的两个张量，则谁先被用到谁排前面。
   排序的结果是一个由张量索引值组成的列表，后续将按照这个列表为张量分配内存。前面说到的ordered_allocs_里面元素的顺序有大致这么确定，不过可能会略微有区别（当存在一个张量太小，可以插入其他一些张量的空隙，这点在后面有介绍）。

3. 接着，根据第二步得到的列表依次为张量分配内存偏移地址，分配过程如下所示，整个过程很简单，从头到尾依次查看是否有空隙容纳当前张量，如果没有则再已分配的张量后面进行分配。

你肯定好奇，既然每个张量的起始地址都进行了内存对齐，那么哪来的空间进行插入操作？岂不是脱裤子放那啥——多此一举。别急，下面就是见证奇迹的时刻：

TfLiteStatus SimpleMemoryArena::Allocate(
    TfLiteContext* context, size_t alignment, size_t size, int32_t tensor,
    int32_t first_node, int32_t last_node,
    ArenaAllocWithUsageInterval* new_alloc) {
  TF_LITE_ENSURE(context, alignment <= arena_alignment_);
  new_alloc->tensor = tensor;
  new_alloc->first_node = first_node;
  new_alloc->last_node = last_node;
  new_alloc->size = size;
  if (size == 0) {
    new_alloc->offset = 0;
    return kTfLiteOk;
  }

  // If we don't find a better gap just allocate at the end of the buffer.
  const size_t kOffsetNotAssigned = std::numeric_limits<size_t>::max();
  size_t best_offset = kOffsetNotAssigned;
  size_t best_offset_fit = kOffsetNotAssigned;

  // Go through the sorted allocs and look at the gaps between them.
  size_t current_offset = 0;
  for (const auto& alloc : ordered_allocs_) {
    if (alloc.last_node < first_node || alloc.first_node > last_node) {
      // Usage interval of alloc doesn't intersect with current tensor's usage
      // interval, so we skip it.
      continue;
    }
    size_t aligned_current_offset = AlignTo(alignment, current_offset);
    // If we found a gap larger than required size, and smaller than previous
    // best fit, take it.
    if (aligned_current_offset + size <= alloc.offset &&
        alloc.offset - aligned_current_offset < best_offset_fit) {
      best_offset = aligned_current_offset;
      best_offset_fit = alloc.offset - current_offset;
    }
    current_offset = std::max(current_offset, alloc.offset + alloc.size);
  }
  if (best_offset == kOffsetNotAssigned) {
    best_offset = AlignTo(alignment, current_offset);
  }

  // Update the required buffer size.
  high_water_mark_ = std::max(high_water_mark_, best_offset + size);
  new_alloc->offset = best_offset;

  auto insertion_it = ordered_allocs_.begin();
  while (insertion_it != ordered_allocs_.end() && *insertion_it < *new_alloc) {
    ++insertion_it;
  }
  ordered_allocs_.insert(insertion_it, *new_alloc);
  return kTfLiteOk;
}

上面展示的是进行内存分配的核心代码——Allocate()，通过上面的代码我们知道，如果不需要保存中间结果，那么TFLite用于张量分配的空间大小为输入张量所需空间之和 + 输出张量所需空间之和 + 节点所用张量中最大的张量所用空间 * (1 + 最大分支数)。假设所有张量大小都为64字节（为了计算方便），模型有两个输入张量、一个输出张量、十个中间结果张量、不包含分支，如果不需要保留中间结果，那么所需空寂大小为（2 + 1 ）* 64 + 64 * 2 = 320 bytes，与之对比的是保留中间结果所需的内存（2 + 1 + 10）* 64 = 832 bytes，可以看到节约的内存还是相当可观的，减少了将近2/3的内存使用。

那么，TFLite是怎么做到的呢？
还记得我们前面我们前面说过的alloc_node_和dealloc_node_吗？他们的用处就在这体现出来了。Allocate()调用时传递的参数如下所示。

          arena_.Allocate(context_, tensor_alignment_, tensor.bytes,
                          tensor_index, alloc_node_[tensor_index],
                          dealloc_node_[tensor_index], &allocs_[tensor_index]));

结合上面的Allocate()的源码，我们举个栗子：

假设有一个模型，一共有6和节点编号0~6，用圆圈，以及五个用于保存中间结果Tensor，编号i~v，用方形表示。如图4所示，假设这五个张量的大小关系为 i > v > ii > iv > iii。现在我们要为所有张量分配偏移地址。

从前面的描述我们知道，此时这些张量的索引值已经有序的被保存在一个列表里，其顺序为：输入输出张量的索引, i, v, ii, iv, iii，alloc_node_以及dealloc_node_列表的内容如图5所示，输入输出的值前面已经提到过，分别是0 和 0xFFFFFFFF，就不再图中画出。

如下图6所示，蓝色的是输入输出张量所占据的内存，从alloc_node_以及dealloc_node_的值可知，他们之间都是有交集的，因此他们所占的内存会一个接着一个。虚线开始，就是需要为这五个中间结果张量分配的空间。

根据排序结果，输入输出的内存分配完成后紧接着的就是分配张量 i。由于i与所有输入输出张量都有交集，因此只能排在他们后面。和自然的，i的偏移地址就从输入输出张量结束的地方开始。

紧接着，需要分配张量v的内存，同理，它会排在所有输入输出后面，目前它和i的偏移地址是一样的。接下来，需要拿v的偏移地址与所有在它之前已分配好了偏移地址并且与它有交集的张量比较以便进行偏移地址的调整。很显然，目前只有i分配完成，但是从alloc_node_以及dealloc_node_的值来看，i的范围是1~2，而v的范围是4~5，和它没有交集，因此偏移地址不用调整，因此v的偏移地址会和i的一样。

同样的，接下来分配的是ii的偏移地址，由于ii与i、v有都有交集，因此需要根据他们调整偏移地址，由于i所需空间大于ii，因此在有着相同的起始地址情况下，肯定是i往后延伸跟多，因此ii直接回跟在i结束的地方。

接下来，分配iv的偏移地址。iv只与v有交集，因此只根据v进行调整，所iv紧随v的后面。

最后，分配iii。iii与ii、v、iv都有交集，因此一共会调整三次，最终会接在ii或者iv的结尾，就看谁结束的地方更靠后。我们假设最终会接在ii结束的地方。如图，最终他们占据的内存最多为i + ii + iii。v与iv复用同一片地址空间。

我们可以推演下整个推过程：

1. 首先i会被赋值；
2. 接着，有i计算得到ii，此时i已经没用了;
3. 然后，不管我们先计算iii还是v，都不会覆盖到ii的值，只会覆盖i；
4. iii、v计算完成后，ii也几经没用因此它的空间又可以为iv所用。

可以看到，整个过程不仅节省了空间，而且也能保证我们的数据都是安全的。如果需要保存中间结果，那么整个内存排布将变成下面这样：

好了今天就到这了。欢1迎2关3注4个5人6微7信8公9众10号【爱码士1024】，此地有崇山峻岭，茂林修竹；又有清流激……

等等！你还没解释为什么内存对齐了还会存在空隙插入比较小的张量！

嗷对！为什么对齐了还会有空隙呢？还是前面的例子，极端点，在i特别大的情况下（Q:请问多大叫特别大？A: 鲲之大，一锅炖不下那种……），当我们分配iv的时候，在iv的结尾到iii的起始这段空间就会出现一个大间隙。如图12所示，这不就可以趁虚而入了。这个故事告诉我们，在特别堵的路上跟车一定要跟的紧，否则肯定会有不讲武德的人插队……

3.3. 实际分配

这个内存分配过程挺巧妙的。完成这一步，距离成功只有一步之遥，最后只要成功向系统申请一块足够大的内存就行了。就相当于我们常说的“等我有了钱，我要怎么怎么样”，计划已经做好了，现在就差两块钱去买体彩。唯一的区别就是操作系统大概率会满足TFLite的需求。

其实到了这一步非常简单了，由于每个张量都有了偏移地址，因此只需要简单的加上申请下来的基地址，就能得到所有真实地址。这里就不多介绍了。

首发于个人微信公众号【爱码士1024】。微信扫描上方二维码或者微信搜索TensorBoy并关注，及时获取更多最新文章！
C++ | Python | 推理引擎 | AI框架源码，有一起玩耍的么？

4. Resources

[1] https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite