MegEngine 动态执行引擎 Imperative Runtime 架构解析

在之前的文章中我们介绍过 MegEngine 的 Imperative Runtime 以及它与 MegBrain、MegDNN 的关系，这篇文章中我们将介绍 Imperative 中包含的常用组件。

在 MegEngine 中，从用户在 python 层编写代码到在 interpreter 层发生计算经过了下面的流程：

用户在 python 层编写网络结构代码，执行时向 C++ 层发射算子执行指令
Imperative 的 dispatcher 对部分算子做计算前的预处理（transformation）
Imperative 的 interpreter 执行计算操作（复用 MegBrain 的相关组件）

我们将分别介绍这几个阶段系统所做的工作。

MegEngine 的 Python 层

在主流的深度学习框架中，用户往往不需要自己手写算子的具体实现、处理计算图的执行逻辑、或者与复杂的体系结构打交道。一切都被封装为 Python 层的接口。

在 MegEngine 的 Python 层中用户接触较多的模块主要有：data、functional、module、optimizer、quantization、tools，下面简单介绍一下各个模块的功能。

构建数据处理 Pipeline —— data 模块

Data 模块，顾名思义就是对数据进行处理的模块。

没有数据就没法训练，在 MegEngine 中，通常会借助一个 Dataset 结构来定义数据集。数据集一般分为 Map-stype 和 Iterable-style 两种。前者叫作 ArrayDataset，这种数据集支持随机访问；后者叫作 StreamDataset，因为是流式的数据集，只支持顺序访问。

有了数据集，我们还需要一个结构来把数据“喂”给模型训练，这样的一个结构叫作 dataloader。

实际上，只给 dataloader 一个 dataset 有时无法准确地描述加载数据的整个过程，我们可能还需要定义加载数据过程的抽样规则（Sampler），或者定义一些数据变换的规则（Transform），或者是定义抽样后的数据的合并策略（Collator）。

Python 层计算接口 —— functional 模块

深度学习模型通常包含一些基础的计算操作，比如 convolution、pooling 等，在 python 层，这些基本计算操作都定义在 functional 模块中。

functional 中实现了各类计算函数，包含对很多 op 的封装，供实现模型时调用。

模型结构的小型封装版本 —— module 模块

使用 functional 提供的接口已经足够编写神经网络模型的代码，但随着模型结构的复杂程度加深，多次反复编写相似的结构会使开发和维护成本迅速提高。

考虑到神经网络模型通常是由各种层（layer）组成，我们通常使用 Module 来封装模型的部分结构或者层，用户实现算法时往往使用组合 Module 的方式搭建模型计算的 pipeline。定义神经网络时有些结构经常在模型中反复使用，将这样的结构封装为一个 Module，既可以减少重复代码也降低了复杂模型编码的难度。

使用 optimizer 模块优化参数

MegEngine 的 optimizer 模块中实现了大量的优化算法，同时为用户提供了包括 SGD、 Adam 在内的常见优化器实现。这些优化器能够基于参数的梯度信息，按照算法所定义的策略对参数执行更新。

降低模型内存占用利器 —— quantization 模块

量化是一种对深度学习模型参数进行压缩以降低计算量的技术。它基于这样一种思想：神经网络是一个近似计算模型，不需要其中每个计算过程的绝对的精确。因此在某些情况下可以把需要较多比特存储的模型参数转为使用较少比特存储，而不影响模型的精度。

MegEngine 相关工具汇总 —— tools 模块

用户进行开发时有时需要一些工具进行错误调试或者性能调优，tools 下就提供了一些这样的工具。比如对训练程序进行记录并在浏览器上可视化的 profiler、方便用户查看 MegEngine 显存占用的 svg_viewer 等。

一般来说，用户会基于上面的模块搭建算法模型，其中定义了非常多的 op 的计算过程，下面我们看一下 c++ 是怎么进行这些 op 的真正的计算的。

Dispatcher 会对 op 做哪些处理？

从 Python 层往下的部分用户往往是感知不到的，脱离了“前端”，我们抽丝剥茧，进入到了框架“后端”对 tensor 和 op 处理的细节。

前面我们提到在 functional 模块中封装了很多算子，并以 python 接口的形式提供。实际上这些算子需要向下发射指令对 tensor 进行操作并返回操作完成后的 tensor，这些发射的 op 指令就会到 dispatch 层，在进行实际计算之前，dispatcher 会对 tensor 做一些处理，我们把这些处理叫作 Transformation。

在 imperative 中真正执行算子进行计算是在 interpreter 层做的，与 tensor 处理相关的操作被解耦出来放在 dispatch 层，这样更便于维护。

在 MegEngine 中，一些重要的 transformation 有：

DimExpansionTransformation：某些 op 计算时对输入 tensor 的 shape 有要求，在这里做处理。
DtypePromoteTransformation：某些 op 要求计算的 tensor 拥有相同的类型，会将所有的输入的类型提升为同一类型之后再进行计算。比如 int 类型 tensor 和 float 类型 tensor 进行计算，需要把 int 类型的 tensor 转换为 float 类型 tensor。
InterpreterTransformation：顾名思义，这类 Transformation 将指令转发到 Interpreter 层（Interpreter 可以认为是 Imperative 中所有计算操作的入口）进行计算，并获取指令的计算结果。Transformation 通常是叠加的，InterpreterTransformation 是最后一层，其后不再跟其他的 Transformation 处理。
FormatTransformation：由于在不同情况下对不同 format 的 Tensor 的计算速度不同，因此需要对 NHWC 和 NCHW 的 Tensor 进行转换，为了不让用户感知到这样的转换，这部分的工作由 FormatTransformation 完成。
GradTransformation：训练模型时需要通过反向传播更新模型参数，反向传播需要支持 op 的自动微分。要实现求导，就需要在前向执行 op 的时候记录某些信息，以便之后进行反向求导。Autodiff 算法会根据输入的前向图生成一个完整的前向反向图，所谓的前传反传训练过程对 Autodiff 来说实际上都是一个计算图的前向过程，grad 的数值是在“前向”的过程中就已经拿到的。GradTransformation 处理的就是与反向求导相关的操作。
TracingTransformation：
在介绍 Trace 之前，我们需要先明确一下计算图的概念。计算图可以认为是对输入的数据（tensor）、op 以及 op 执行的顺序的表示。计算图分为动态图和静态图。动态图是在前向过程中创建、反向过程销毁的。前向逻辑本身是可变的，所以执行流程也是可变的（因此叫动态图），而静态图的执行流程是固定的。也就是说，动态图在底层是没有严格的图的概念的（或者说这个图本身一直随执行流程变化）。对于动态图来说，graph 的 node 对应的概念是 function / 算子，而 edge 对应的概念是 tensor，所以在图中需要记录的是 graph 中 node 和 edge 之间的连接关系，以及 tensor 是 function 的第几个输入参数。
Trace 的作用就是将动态图执行转换为静态图执行，这样做的好处就是执行速度更快了，并且占用的显存更少了。因为静态图需要先构建再运行，可以在运行前对图结构进行优化（融合算子、常数折叠等），而且只需要构建一次（除非图结构发生变化）。而动态图是在运行时构建的，既不好优化还会占用较多显存。
Trace 中所有的东西都会进行静态优化（加速）。
加了 Trace 之后，模型在训练时第一个 iter 是动态图执行，Trace 会记录下 tensor、op 以及 op 的执行顺序这些信息（构建静态图）并进行计算，在第二个 iter 就跑的是构建好的静态图。
LazyEvalTransformation：类似 TracingTransformation，也会记录 tensor、op 等信息构建静态图，不同的是 LazyEvalTransformation 在第一个 iter 不会跑动态图，但会在第二个 iter 开始跑静态图。
ScalarTransformation：用于判断指令的输出是否为 scalar。因为 dispatch 的 Tensor 要发到 Interpreter 层，而 Interpreter 层不接受 ndim == 0 的 Tensor（在 Interpreter 中 ndim 为 0 表示 Tensor 的 shape 未知），也就是一个 scalar，因此 ScalarTransformation 会将 ndim 为 0 的 Tensor 表示为 ndim 不为 0 的 Tensor （具体是多少与具体 op 有关）发往 Interpreter。

不同的 Transformation 之间拥有固定的执行顺序：比如 InterpreterTransformation 是执行实际计算并获取计算结果的（需要进入 Interpreter），所以它是在最后一个执行的。TracingTransformation / LazyEvalTransformation / CompiledTransformation 等属于 Trace 相关的操作，因为 Trace 需要记录所有指令，所以这些 Transformation 是在倒数第二层执行的。如 ScalarTransformation 这样只对 Scalar 做处理的 Transformation 往往在较上层。

因为不同的 Transformation 有逻辑上的先后关系，所以开发者往往需要手动规划它们之间的顺序。

不同类型的 Transformation 之间是解耦的，这样便于开发与维护。

Interpreter 是如何“解释”算子的？

由于 MegBrain 已经是一个非常成熟的静态图框架，因此在开发动态图（Imperative Runtime）深度学习框架 MegEngine 的过程中，复用许多静态图中的组件可以大大降低开发成本。

实际上，张量解释器 Tensor Interpreter 就是将动态图中的操作——如执行 op、shape 推导等操作“解释”为静态图的对应操作，并复用 MegBrain 的组件来运行。

这里我们需要先了解一个 MegBrain 的静态图“长什么样”。

复用静态图接口的机制 —— proxy_graph

为了复用 MegBrain 的静态求导器、静态内存分配器、静态 shape 推导器等组件，imperative 引入了 proxy_graph。

复用 MegBrain 的接口需要实现对应的方法，在 MegEngine/imperative/src/include/megbrain/imperative 目录下可以看到所有需要实现的桥接接口，其中和 proxy_graph 相关的接口声明在 proxy_graph_detail.h 中，通常需要实现这几个接口：

infer_output_attrs_fallible
复用 MegBrain 的 StaticInferManager 进行 shape 推导，在执行计算操作前对输入和输出 tensor 的 shape 进行检查。
apply_on_physical_tensor
根据 infer_output_attrs_fallible 推导的 shape 结果去分配 op 输出的显存，并调用 proxy opr 的 execute 函数（会转发到 MegDNN 的 exec 函数）执行计算操作。
make_backward_graph
在求导时，Grad Manager 会记录下来一些求导需要的信息（输入 tensor、op 以及它们执行的顺序、输出 tensor），make_backward_graph 会根据这些信息造一个反向的计算图，供求导使用。
get_input_layout_constraint
一般用来判断一个输入 tensor 的 layout 是否满足一些限制：比如判断 tensor 是否是连续的。
如果不满足限制，则会造一个满足限制的 tensor，供 apply_on_physical_tensor 使用。

在实现一个 imperative 算子时通常也只需要实现这几个接口，剩下的工作由 MegBrain 和 MegDNN 完成。

主流框架在 python 层的模块封装结构大同小异，关于 MegEngine 的 Python 层各模块的使用与实现细节以及 transformation 和 interpreter 实现细节我们会在之后的文章中逐一解析。

附

更多 MegEngine 信息获取，您可以：查看文档和 GitHub 项目，或加入 MegEngine 用户交流 QQ 群：1029741705。欢迎参与 MegEngine 社区贡献，成为 Awesome MegEngineer，荣誉证书、定制礼品享不停。