当我们谈论 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
。这样做的原因是:
- 重写算子代价高,且容易写错。
- 若
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
。