计算图的生成

主流机器学习框架都支持动态图和静态图。

动态图是即时编译的,也就是有一个 op 和输入之后立即发 kernel 执行计算得到结果并返回。

静态图是获取到整个程序结构之后构建了计算序列之后再进行计算的,比如有 op1、op2…. opn,还有输入 input1、input2…inputm,然后执行整个拓扑序列的计算。

从动态图的特点能很容易的发现动态图的优点就是可以很方便地对于单个 op 得到结果、方便调试。缺点就是没有获取全局信息没法做很多优化(比如算子融合,因为还不知道下一个 op 是什么就已经发 kernel 执行了)。

从静态图的特点看可以发现静态图的优点就是掌握了全局信息可以做更多的图优化(算子融合等),这样可以让我一个计算过程的总时间更快(因为可能发了更少的 kernel,dispatcher 到后端的通信也更少了),缺点就是调试不方便,比如我只想调试一两个 op 的结果却不得不跑整个计算图。

在 PyTorch 和 MegEngine 中默认跑的是动态图,如果要跑静态图需要给方法加上 trace 装饰器把动态图转成静态图(这种情况下在第一个循环跑的是动态图,框架会记录下拓扑结构,在第二个循环开始就跑的是静态图了)。

静态图的生成

静态图的生成采用先编译后执行的方式,也就是我刚才说的记录下拓扑结构的方式,将计算图的定

义和执行分离:

使用前端语言(Python)定义模型形成完整的程序表达后,mlsys 首先对模型进行分析,获取网络层之间的拓扑关系(有哪些 op、每个 op 的输入输出结构、op 之间的计算顺序)还有参数变量设置、损失函数等信息。然后 mlsys 会将完整的模型描述编译为可被计算后端调用执行的固定代码文本(或者叫 IR,intermediate representation 中间表示),这种有模型完整拓扑关系的计算图就是静态图。静态计算图直接接收数据并通过相应硬件调度执行图中的算子来完成任务。静态计算图可以通过优化策略转换成等价的更加高效的结构,提高后端硬件的计算效率(也就是对 IR 做一些优化)。

举个例子,下面有一个程序:

def model(X, flag):
	  if flag>0:
        Y = matmul(W1, X)
	  else:
        Y = matmul(W2, X)
    Y = Y + b
    Y = relu(Y)
return Y

这里第二行有条件判断,在实际执行的时候可能执行不同的计算(W1 或者 W2 和输入 X 计算 matmul),因为静态生成模型没有数据输入,所以我们需要将条件控制算子以及所有的分支计算子图加入计算图中。在静态计算图执行计算阶段网络接收数据流入,调度条件控制算子根据输入数据进行逻辑判断,控制数据流入不同的分支计算子图中进行后续计算。在部分机器学习框架中前端语言Python的控制流不能够被正确编译为等价的静态图结构,因此需要机器学习框架的控制原语来实现控制流。

上图中将静态计算图的结构转化为一个等价的计算图,走了某个分支的计算子图并且还做了算子的融合,然后再把这个等价计算图发到后端去做计算操作。

除了掌握全局信息便于优化以外,静态图还有一个优势就是可以做模型部署。在模型推理阶段,执行序列化的模型即可,无需重新编译前端语言源代码。机器学习框架可以将静态计算图转换为支持不同计算硬件直接调用的代码。结合计算图序列化和计算图转硬件代码两种特性,静态图模型可以直接部署在不同的硬件上面,提供高效的推理服务。

静态图最大的缺点就是不便于调试:比如在代码中,若add算子和relu算子经给优化合并为一个算子,执行时合并算子报错,用户可能并不知道错误指向的是add算子错误 还是relu算子错误。此外在神经网络模型开发迭代环节,不能即时打印中间结果。若在源码中添加输出中间结果的代码,则需要将源码重新编译后,再调用执行器才能获取相关信息,降低了代码调试效率。

动态图的生成

动态图的特点是编译和执行同时进行(jit,just-in-time compilation,即时编译),有了一个 op 立即发 kernel 计算结果并返回。

比如刚才的那个静态图的例子,每个 op 都会立即计算:

反向的时候会按照和前向相反的顺序发对应的反向 kernel 也做一遍计算。

前向计算 matmul(w1, x) 的时候,会记录下来反向计算过程是 grad_w1 = grad_y * x,记录下需要参与反向计算的算子和张量 x,mlsys 依据收集的信息完成前向计算和反向图构建。

尽管动态生成中完整的网络结构在执行前是未知的,不能使用静态图中的图优化技术来提高计算执行性能。但其即刻算子调用与计算的能力,使得模型代码在运行的时候,每执行一句就会立即进行运算并会返回具体的值,方便开发者在模型构建优化过程中进行错误分析、结果查看等调试工作。

动态图和静态图对比

特性 静态图 动态图
即时获取中间结果
代码调试难易
控制流实现方式 特定的语法 前端语言语法
性能 优化策略多,性能更佳 图优化受限,性能较差
内存占用 内存占用少 内存占用相对较多
是否可直接部署 可直接部署 不可直接部署

参考

posted @ 2023-08-05 23:21  machine_gun_lin  阅读(171)  评论(0编辑  收藏  举报