【NNVM 记录】计算图的内存分配
计算图是深度学习中最基础也最重要的一种数据结构,在计算图的运行过程中,每个节点会依次执行具体的算子操作。此时,需要有一定的内存空间来存放算子的输入和输出。这篇文档分析了 NNVM 中计算图内存分配部分的具体处理过程。
编译计算图
编译计算图的过程由GraphCompile
执行,该 pass 位于nnvm/src/compiler/graph_compile.cc
。
作为在 build 的最后阶段被调用的 pass,它的处理流程是:
- 处理所有可融合的算子节点,然后将其转化为融合后的节点,包括将融合后的算子编译成为
LoweredFunc
; - 建立原图和融合后的计算图的映射关系,并且据此构建出新的计算图。这一过程中会对
assign
算子进行特殊处理:标记所有的assign
节点,并且检测可以简化的assign
,将其特殊标记并转为空操作; - 将处理完成的计算图进一步 build 为 module,得到目标代码;
- 执行内存分配,并且对 Placeholder、
assign
节点的内存分配情况作出额外处理。
在这篇文档将讨论 NNVM 的内存分配机制,所以我们将主要关注上述第二步和第四步。