作者：周瑞亮 | 旷视科技 MegEngine 架构师

在正式进入本文前，先通过一组简略的比照来直观地感受一下，动态内存优化带来的内存占用上的晋升。别离对 1batch/128batch 的 vgg16/resnet50/mobilenetv2 模型，应用 tflite tools 失去 tflite 的峰值内存，与 megengine 内存优化前后的峰值内存进行比照：


从上图能够看出，通过动态内存优化，模型占用的内存空间显著减小。接下来，将为大家揭秘 MegEngine 深度学习框架的动态内存是如何治理优化的。

1、升高内存占用——为什么以及怎么做

why

近几年深度学习都在蓬勃发展，并被利用于各行各业。随着业务场景对精度要求的一直进步，深度学习模型的规模也越来越大。 宽广的深度学习钻研人员无论是应用何种运算平台，内存的占用都是绕不开的话题：内存的占用量间接关系到可训练模型的规模，以及部署时的性能体现。鉴于内存的大小无限，内存治理成为深度学习零碎的重要研究课题。

how

内存治理是深度学习零碎的重要研究课题。针对这个问题外界曾经存在许多计划，典型计划如下：

  • 模型压缩，例如量化办法、模型剪枝等。
  • 计算换内存：例如往期分享的 DTR 技术，抛弃计算成本绝对较低的值，并在反向流传时从新计算它们。
  • 通信换内存：将未应用的变量从以后运算设施的内存空间替换到其它内存空间，并在下次访问之前将它们替换回来，会产生通信和同步的开销。
  • 对顺序程序图中的数据流进行剖析，以容许重用内存。

2、MegEngine 中动态内存治理

MegEngine 中采纳了多种升高内存的解决方案，本次次要介绍的是 MegEngine 动态内存治理模块是如何利用顺序程序图中的数据流剖析，实现内存重用以达到升高内存占用的成果。MegEngine 的思路是在运行前计算失去所有的动态内存总共须要的内存池的大小，而后为每个动态内存申请指定其在内存池中的偏移地址，运行时应用曾经预先指定好的空间即可。所以，整个算法的外围在于如何获取总的内存池大小以及为每个动态内存申请调配偏移地址。该算法大抵能够分为三步：

  • 获取原始数据流信息，从中解析失去动态内存的需要信息、依赖关系等。
  • 将动态内存的需要信息形象为 memory chunk，通过依赖关系对 memory chunk 进行初步优化，将繁琐的动态内存需要信息缩减为老练的 memory chunk。
  • 依照依赖关系对 memory chunk 的申请流程进行仿真实现，依据仿真后果失去总的内存池大小以及为每个 memory chunk 调配偏移地址。

从上述三步能够看出 MegEngine 的动态内存 接下来具体拆解这三步的实现办法。

a、获取原始数据流信息

MegEngine 会将图通过编译失去一个拓扑排序的算子执行序列，在运行时会依照算子的执行序列程序执行。通过这个算子序列，以及这些算子的输出/输入节点的信息，能够轻松地察看到在整个模型运行时，数据是如何被传输应用的。这便是一个残缺的顺序程序图的数据流信息。

b、建设最小内存治理单元 Memory Chunk

依据顺序程序图的数据流信息，通过内存治理模块，建设最小内存治理单元 memory_chunk，失去它们的读/写依赖关系以及生命周期。

（1）遍历算子（opr）执行序列，并依据定义在算子中的办法来推断该算子输出（input）/输入 (output) 节点之间依赖关系，根据依赖关系，为每个输出/输入节点创立 memory chunk：

(readonly 和 writable 地址关系示意图）

如上图所示，input0/output0 是读依赖，input1/output1 是写依赖，它们在地址上的映射关系如图中右侧所示，每个矩形对应于一个 memory chunk 的概念。memory chunk 之间的依赖关系以及创立规定如下：

  • 读依赖关系，要求某个输入节点复用输出节点的内存，并且保障不会批改内容。这种状况下，输入节点不创立新的 memory chunk, 而是间接复用输出节点的 memory chunk。如图所示，input0 和 output0 这两个节点共用一个 memory chunk 即 chunk0。
  • 写依赖关系，输入节点会复用输出节点的全副或局部内存。这种状况下，输入节点会创立新的 memory chunk(chunk2)，并记录写操作执行后的 memory chunk(chunk2) 在原先 memory chunk(chunk1) 中的偏移地址以及占用的空间大小。
  • 和之前的输出都没有关系，独自创立一个 memory chunk，记录占用空间的大小等信息。

（2）获取每个 meory chunk 的生命周期：

  • 先将所有 memory chunk 的起始生命周期（begin）和终止生命周期（end）都记为 0。

  • 程序遍历算子执行序列

    • 获取输入节点，如果该输入节点首次被统计，那么该节点的 memory chunk 的起始生命周期记为算子的执行序号 id，意味着该 memory chunk 的数据在第 id 个算子执行后产生；
    • 获取输出节点，将该节点的 memory chunk 的终止生命周期（end）和 id+1 比照，并将 end 更新为两者中的较大值，即 end=max(end, id+1)。id+1 意味着要在第 id 个 opr 执行完结后能力开释。

须要留神的是，在上述步骤中“首次统计”以及“max(end, id+1)”的形容都是因为读依赖关系导致存在多个输出/输入节点是共用一个 memory chunk，该 memoy chunk 会被屡次统计到，利用算子执行序列的拓扑性质以及上述办法来确保生命周期统计的正确性。

在获取每个 meory chunk 的生命周期的步骤中，读依赖的性质曾经被用于内存优化：它一直地缩短该 memory chunk 的生命周期，以缩小不必要的内存操作。

c、内存优化算法对 Memory Chunk 进行优化

接下来，内存优化算法会利用上文 b(1) 处失去的 memory chunk 之间的写依赖的性质以及 b(2) 处失去的 memory chunk 的生命周期进行进一步的优化。次要分为两步：

  • 依照依赖关系对 memory chunk 的申请流程进行仿真实现，获取到每个 memory chunk 申请时，前置空间的 memory chunk。
  • 依据仿真实现返回的后果，推算失去总的内存池大小以及为每个 memory chunk 调配偏移地址。

（1）memory chunk 申请空间的仿真实现

首先，介绍算法中内存治理办法以及 free()/alloc()/alloc_overwrite() 的逻辑。

为了模仿内存的调度过程，保护了两个内存信息队列：

  • 全副内存信息队列（后简称内存队列）：按地址程序记录所有曾经申请的内存，包含应用中（灰色局部标记）和闲暇（红色局部标记）两种状态。
  • 闲暇内存信息队列（后简称闲暇内存队列）：按空间大小从小到大程序记录的处于闲暇状态的内存。

保护的准则是：内存队列中不存在间断的闲暇内存，因为在模仿调度过程中，会将间断闲暇内存合并为一个大的闲暇内存，不便后续的调度，缩小不必要的内存碎片的产生。


（内存治理链表）

free() 逻辑：开释指定地址空间，并保护内存队列和闲暇内存队列。次要有两种状况，如下图所示：

  • 状况一：free 之后的闲暇空间前后无其它闲暇空间，不须要合并闲暇空间的操作。
  • 状况二：free 之后的闲暇空间前后有其它闲暇空间存在，须要进行合并闲暇空间的操作。

)

alloc() 逻辑：申请 size 大小的空间，并保护内存队列和闲暇内存队列。次要有三种状况，如下图所示：

  • 状况一：闲暇内存队列为空，则间接申请一块 size 大小的内存，并记录在内存队列开端；
  • 状况二：闲暇内存队列不为空，且有大于或等于 size 的空间，则从闲暇内存队列中申请一块大于或等于 size 的最小空间。截取所需空间，并保护内存队列和闲暇内存队列。

  • 状况三：闲暇内存队列不为空，然而所用空间的大小均小于 size，则取出闲暇内存队列中最大的闲暇空间（即队尾记录的内存）并扩大至 size 大小。

alloc() 胜利后，记录前置空间的 memory chunk（即前一个非闲暇块的 memory chunk）。

alloc_overwrite() 逻辑：上文 b(1) 处提到写依赖关系，能够看到写依赖的节点可能只应用写依赖对象节点原空间的一部分空间（也可能是全副，基于建设写依赖关系时记录的偏移地址和大小信息能够推导）。所以进行写操作之后，前后可能会有闲暇空间产生，将这部分空间开释，并保护内存队列和闲暇内存队列。胜利之后，记录前置空间的 memory chunk。

举一个例子进行阐明，假如有一个写依赖关系如 b(1) 图示：chunk2 对 chunk1 进行写操作，且 offset = 6M，size = 10M，chunk1 的大小为 20M。由此能够推断出，在写操作实现后，原 chunk1 的 20M 空间被合成为后面 6M 的闲暇空间，两头 10M 的写笼罩的空间，前面 4M 的闲暇空间。那么将前后闲暇空间开释的操作就如下图所示：

接下来，依据 b(2) 处失去的 memory chunk 的生命周期，创立 time=>(alloc, free) 的数据结构（time 为 opr 序列，alloc 为该工夫点 alloc 的 memory chunk 的指针数组，free 为该工夫点 free 的 memory chunk 的指针数组），并根据 b(1) 处取得的写依赖的性质，模仿整个内存治理的流程：


（算法模仿 alloc 和 free 流程图）

如上图所示，整个流程能够划分为三步，也对应之前提到的三种办法：

  • 将该工夫点须要开释的且不被写依赖的 memory chunk 调用 free() 办法进行开释。
  • 将该工夫点须要申请的且不是写依赖行为的 memory chunk 调用 alloc() 办法进行申请。
  • 在上述两步实现后，将具备写依赖行为的 memory chunk 调用 alloc_overwrite() 进行解决。

(2) 推导总的内存池大小以及为每个 memory chunk 调配偏移地址

为了便于了解 memory chunk 的偏移地址的记录形式。对上述曾经提及的内容进行一个专门的汇总阐明，所有的 memory chunk 能够分为两类：通过 alloc() 申请的没有写依赖行为的 memory chunk（也称为 memory_chunk_root）和通过 alloc_overwrite() 申请的有写依赖行为 memory chunk。memory chunk 能够根据写依赖关系建设一组链表，如下图，以其中一个链表为例（memory_chunk_n 写依赖于 memroy_chunk_n-1,...,memory_chunk_1 写依赖于 memory_chunk_root，这就形成一个写依赖链表，剖析所有的 memory chunk 写依赖能够失去多个写依赖链表）。因为写依赖关系是会记录 offset 和 size 信息，那么链表中的每个 memory_chunk 都能够推导出与 memory_chunk_root 的 offset_root。那么对于这整个链表上的 memory chunk 只须要调整 memory_chunk_root 的起始地址 (root_addr_begin)，就能够实现对整个链表上 memory chunk 地址的调整。


（overwrite 链表与内存关系）

接下来，介绍如何推导所有的 memory chunk 的偏移地址。依据之前仿真记录的前置空间的 memory chunk 失去前置空间的大小。通过前置空间大小，对 memory chunk 的调配空间进行调整，失去 memory chunk 的偏移地址。 在理论实现中存在递归调用，应用 map 缩小了反复了查问和计算。这里提及的前置空间即仿真记录的前置空间的 memory chunk 以及该 memroy chunk 的迭代推导出的所有前置空间的 memory chunk 占用空间的总和，形象地了解就是调配空间时，地址空间上从 0 开始后方曾经被占用的空间。

如下图所示，纵向示意逻辑地址空间，左侧红色区域示意的是前置空间占用的地址空间，右侧 memory_chunk_n 则是此次须要进行调整空间安顿的 memory_chunk。root_addr_begin 的初始值为 0，调整时有两种状况：

  • 状况一：前置空间已占用大小 <= root_addr_begin + offset_root，这意味着 memory_chunk_n 理论应用空间和前置空间不会产生抵触，所以 root_addr_begin 并不需要调整。

  • 状况二：前置空间已占用大小 > root_addr_begin + offset_root，这意味着 memory_chunk_n 理论应用空间和前置空间会产生抵触，所以 root_addr_begin 须要调整至 root_addr_begin + offset_root = 前置空间已占用大小，以防止抵触同时不产生空间的节约。


（逻辑地址调配示意图）

在此轮调度过后，能够获取到总的内存池大小以及为每个 memory chunk 调配偏移地址：

  • 总的内存池大小：所有 memory_chunk_root 中，root_addr_begin 加上本身大小所达到的最大值。
  • memory chunk 的偏移地址：root_addr_begin + offset_root

在理论运行时，只须要申请一块峰值显存大小的内存，并根据各个 memory chunk 的偏移地址进行显存空间的调度调配，就能够保障程序的失常运行。

3、总结

最初，对 MegEngine 采纳“顺序程序图中的数据流剖析，以容许重用内存”的办法进行内存优化的长处进行一个总结：

  • 能够升高模型的显存占用量。
  • 对于模型齐全通明，不会对模型精度产生影响。
  • 是纯正的内存方向的优化，不会带来计算/通信上的损失。
  • 内存优化是在模型运行前做的，内存治理的运行时开销简直为零。
  • 能提前获取到动态内存的峰值显存，在肯定水平上，防止程序运行到一半因为显存不够导致训练失败浪费时间的状况。
  • 对立的内存申请，升高存储设备上的内存治理老本。

附：
GitHub：MegEngine 天元
官网：MegEngine-深度学习，简略开发
欢送退出 MegEngine 技术交换 QQ 群：1029741705