在之前的文章中我们介绍过 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
实现细节我们会在之后的文章中逐一解析。