乐趣区

关于开源:MegEngine-动态执行引擎-Imperative-Runtime-架构解析

jiezi

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

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

  1. 用户在 python 层编写网络结构代码，执行时向 C++ 层发射算子执行指令
  2. Imperative 的 dispatcher 对局部算子做计算前的预处理（transformation）
  3. 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，荣誉证书、定制礼品享不停。

退出移动版