MegEngine 动态执行引擎 -Imperative Runtime 概述

当我们谈论 MegEngine 时，我们在谈论什么

众所周知，开源深度学习框架旷视天元（MegEngine）是旷视自主研发的国产工业级深度学习框架 ，是旷视新一代 AI 生产力平台 Brain++ 的最核心组件，在 2020 年 3 月正式向全球开发者开源。MegEngine 凭借其训练推理一体、超低硬件门槛和全平台高效推理 3 大核心优势，帮助企业与开发者大幅节省产品从实验室原型到工业部署的流程，真正实现小时级的转化能力。

MegEngine 自上向下包含三个层次：Imperative、MegBrain 和 MegDNN。它们的角色定位分别是：

• Imperative：MegEngine 为动态训练设计的一套新接口，负责处理动态图运行时（Imperative Runtime）。

• MegBrain：负责处理静态图运行时（Graph Runtime）。

• MegDNN：MegEngine 的底层计算引擎。

下面再从下到上介绍一下这三个模块在 MegEngine 中的作用。

MegEngine 的底层计算引擎 —— MegDNN

MegDNN 在 MegEngine 中扮演的角色是底层计算引擎。MegDNN 是一个跨平台的底层算子库，训练和推理过程中的所有计算操作最终都需要落到一个 MegDNN 的算子上进行，一个 MegDNN 算子可能会根据场景（如张量尺寸等）有一个或多个不同的实现（也叫 kernel）。

作为一个跨平台的计算库，MegDNN 为我们提供丰富的与 Tensor 相关的计算原语，比如 Convolution、Pooling、MatrixMul、Transpose 等。目前 MegDNN 支持的平台有：x86、arm、CUDA、RoCM、OpenCL、Hexagon 等。

感兴趣的同学可以在这里看到 MegDNN 在各个平台下的 kernel 具体实现。

静态图运行时管家 —— MegBrain

为了确保训练推理一致性， Imperative 中复用了 MegBrain 的计算代码，因此我们需要了解 MegBrain 做了什么。

MegBrain 负责处理静态图的运行时，主要提供 C++ 的训练和推理接口。

从下面的 MegEngine 整体架构图可以看出，Imperative 通过 Tensor Interpreter （张量解释器）复用了许多 MegBrain 的代码。比如 shape 推导、计算、求导、Trace 等。

在 MegBrain 中，一个 Computing Graph 由 SymbolVar 以及许多 op （算子，operator）组成。SymbolVar 是在 MegBrain 层面 Tensor 的表示，可以理解为传递给 op 进行计算的数据。作个类比，op 是类似加减乘除这样的计算操作（在深度学习中常用的有 convolution、pooling 等），SymbolVar 就是我们用来进行加减乘除的“数”（在深度学习中就是 Tensor）。

MegEngine 动态图接口 —— Imperative 登场

因为 MegEngine 是动静合一的深度学习框架，MegBrain 解决了静态图的训练和推理问题，还需要有一个“组件”负责处理动态图的训练和推理、以及 Python 侧的训练接口，于是便有了 Imperative，也就是说，Imperative Runtime 是为了动态训练而单独设计的一套新接口。

实际上，在 MegBrain 的 Computing Graph 中已经有了非常多的算子实现，因此 MegEngine 的 Imperative 借助张量解释器 Tensor Interpreter 较多地复用了 MegBrain 中的 op。这样做的原因是：

1. 重写算子代价高，且容易写错。

2. 若 Imperative 的实现和 MegBrain 的实现不一致的话，容易导致训练推理不一致。

除了复用 MegBrain 的部分功能，Imperative 自身包含的模块主要有：Module（定义深度学习网络的基础类的实现）、Optimizer（一些优化器的实现）、Functional（提供 python 层的训练接口）、Interpreter（计算入口，底层会调用 kernel 或者 MegBrain 的算子实现）、DTR（动态重计算模块）、Tracer（记录计算图的计算过程）等。

这些模块会在之后的文章中有更详细的介绍，感兴趣的同学可以查阅 MegEngine 官方文档。

总结：Imperative 与 MegDNN / MegBrain 的关系

简单来说，MegDNN 负责 MegEngine 中所有的计算操作在各个平台（CUDA 等）的最终实现，无论是 MegBrain 还是 Imperative 的 op，最终都需要通过调用 MegDNN kernel 来完成计算。

既然 MegDNN 包揽了计算的活儿，那么在训练推理过程中那些与计算无关的工作，自然就落到了 MegBrain 和 Imperative 的头上。这些工作包括：求导、内存分配、对 Tensor 的 shape 进行推导、图优化、编译等。

MegEngine 整体上是有两部分 Runtime 以及底层的一些公共组件组成的。这两部分的 Runtime 分别叫做 Graph Runtime（对应 MegBrain） 和 Imperative Runtime（对应 Imperative）。

Graph Runtime 负责静态图部分，主要提供 C++ 训练推理接口。实际计算时需要调用 MegDNN 的实现。

Imperative Runtime 负责动态图部分，主要为动态训练提供 Python 接口。实际计算时需要调用 MegBrain 的已有实现或者直接调用 MegDNN 的 kernel。

附：

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