OneFlow 如何做静态图的算子对齐任务
撰文|李响
1、前言
深度学习框架中模型的运行方式主要有动态图和静态图两种,动态图更易用,静态图性能更具优势,OneFlow 习惯将它们称为 Eager 模式和 Graph 模式。
OneFlow 提供了 nn.Graph 模块,让用户可以用类似 Eager 模式的编程习惯,构建静态图训练测试。因此,需要同时保证 Eager 和 Graph 模式下算子行为和结果的正确性。
在之前的文章《深度学习框架如何优雅地做算子对齐任务》中 ,分析了 Eager Ops 的自动测试流程,包括如何产生随机数据测试用例和 AutoTest 核心代码实现,AutoTest 框架可以很轻易移植到其它深度学习框架使用。
不过,本文的主要目的则是介绍 OneFlow 如何完成 Graph 模式下算子的测试任务。目前为止,OneFlow v0.7.0 已经新增所有 Op 在 nn.Graph
上做静态执行的单测支持,自动化单测功能完备。
文章中涉及到的代码位置:
https://github.com/Oneflow-Inc/oneflow/blob/master/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object
pyhttps://github.com/Oneflow-Inc/oneflow/blob/master/python/oneflow/test_utils/automated_test_util/generators.py
2、OneFlow 的 Graph 算子对齐概述
OneFlow 提供的 Eager 模式,用法与 PyTorch 对齐。所以在测试上,AutoTest 框架会随机出各种合法参数组合成的 Op ,并基于数值和类型完全相同的输入 Tensor(PyTorch 和 OneFlow 各有一份)分别运行 PyTorch 和 OneFlow 的代码,来完成算子对齐工作。
此外,OneFlow 还提供了 Graph 模式,基于面向对象的编程风格,让熟悉 Eager 开发的用户,只需改很少的代码,就可以高效使用静态图。对比 Eager 模式,Graph 模式不易调试,但性能更好,易于优化和部署。那么,如何自动测试 Graph 模式下的 Ops 就是重点需要关注的问题。
在详细介绍 Graph 单测之前,我们先看一下 AutoTest 框架里 Graph 打开方法,下面是一个测试 matmul 算子的例子。基于 random_pytorch_tensor
方法构造了两个随机的 tensor,它们的维度分别是 [n, k]
和 [k, m]
,这些维度的值都是随机生成的,AutoTest 框架参数的随机性都是基于 generators.py 中的 generator
基类完成的。
通过调用 torch.matmul(x, y)
,自动测试框架会分别运行 Torch 和 OneFlow 的 matmul 算子,会检查 Eager 模式下 OneFlow 和 PyTorch 算子的前向和反向结果是否一致。值得注意的是,代码中 @autotest
装饰器的 check_graph
开关为 True
,表示此时会并行地做 Graph 的单测。
3、Graph 模式下自动测试实现原理
在了解背景和使用方法后,这里介绍 Graph AutoTest 的实现思路。
3.1 AutoTest 流程介绍
在 Eager 的自动测试原理中,关于随机数据是如何产生的和 autotest()
装饰器的实现,在前文中有清晰的介绍。关于 AutoTest 框架核心流程实现,首先必须要关注用于 OneFlow 和 Pytorch 的算子对齐任务中的 GetDualObject
函数。
GetDualObject
函数会重写传入的原始 PyTorch
以及 OneFlow
对象的 __call__
魔法函数,最后返回一个 DualObject
对象。这个过程中还包含跳过一些不需要关注的魔法函数,检查传入对象的属性是否合法,基于 nn.Module
和其它 API 默认参数的类型对 generator
继承类产生的随机数据绑定特定类型的工作( get_args
函数中完成)。此外,在代码中还有对 tensor
方法的特判,因为 tensor
方法的调用方式(通过 getattr
)和 nn.module
、 nn.functional
不同(通过 __call__
)。
基于上述流程,通过执行样例代码中的 torch.matmul(x, y)
,AutoTest 框架会通过调用 GetDualObject
函数生成 DualObject
对象,其中 torch
就可以理解为一个 DualObject
对象。最后执行 DualObject 对象,完成结果对比。自动测试流程中 Eager 算子对齐更多的细节也在前文中做了清晰介绍。
GetDualObject
函数实现:https://github.com/Oneflow-Inc/oneflow/blob/7fe29cb3d24be41fa981c4ad6be3051dacc3b605/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py#L600DualObject
类对象实现在:https://github.com/Oneflow-Inc/oneflow/blob/0826518cc49200dccada0f54d5c83accb9218c83/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py#L784
3.2 Graph 模式如何伴随 Eager 模式做算子对齐
从上面的分析中,可以大概总结出 AutoTest 的流程:生成随机数据、生成 DualObject
对象、执行 DualObject
对象和判断结果是否对齐。其中,在执行 DualObject
对象阶段,AutoTest 框架会并行地执行 OneFlow 算子的 Graph 版本,这样也就完成了 Graph 模式伴随 Eager 模式做算子对齐的任务。此外,本节也梳理了 GetDualObject
函数中应该如何识别需要静态(Graph)执行的对象。
在算子对齐任务中,存在 nn.module
, nn.functional
和 tensor
方法三种类型。这里先以 nn.Module
类型为例,分析 Graph 模式伴随 Eager 模式测试的代码,其他三种类型处理方法基本一致。代码执行顺序如下图:
在 oneflow_eager_run_with_graph_check
中分别调用了 get_module_graph_test 和 get_oneflow_eager_res
,得到 Graph 和 Eager 模式的两个结果,最后检查是否对齐。
也就是说,对于一个测试 case,AutoTest 框架总共执行了 Pytorch、OneFlow Eager 模式和 OneFlow Graph 模式三种代码,来验证三种结果是否都对齐了。
我们先来探究下 get_module_graph_test
这个接口,也就是如何得到 Graph 版本的计算结果。代码如下:
其中 oneflow
是一个 nn.Module
对象,graph_train_oneflow
是它的深拷贝结果,主要是为了防止在测试算子的 inplace 版本时,对相应的 DualObject
对象值进行了修改,造成 Graph 的输入和 Eager 不一致导致测试结果对不齐的情况。
首先为了验证 Graph 的后向可以正常执行,构造了一个 Optimizer。在__init__
中复用 Eager 模式下的 nn.Module
对象后,在 build
中描述了 Graph 测试的计算过程,最终返回了 Graph 的实例。简单来说,就是构造一个适应所有算子的通用静态图模型。
在讨论如何构造静态执行代码计算 Graph 结果之后,识别需要静态执行的对象也是需要优先解决的问题。 oneflow_eager_run_with_graph_check
的完整代码如下:
在 oneflow_eager_run_with_graph_check
中,需要判断哪些对象需要静态执行测试。因为 OneFlow 设定部分代码需要静态化,比如有些 Eager 模式下的方法,在 Graph 模式下没有定义。
上面的代码中首先通过 if testing_graph
: 判断是否打开了 Graph 开关,既是否需要并行的做 Graph 的单测;再对 oneflow
对象的类型做 isinstance
判断,当为 nn.Module
时才需要静态执行,调用 get_module_graph_test
。否则调用 get_functional_graph_res
等处理,在测试框架中其他需要类似判断的地方也同理。
3.3 Graph 模式的自动测试个性化
在 3.2 介绍了 Graph 如何伴随 Eager 模式做算子对齐的任务之后,本节主要分析 Graph 模式自动测试的个性化内容。
在 Graph 模式下,需要处理 nn.module
、nn.functional
和 tensor
三个类别的方法,AutoTest 框架采用先判断后构图的方式。
首先,GetDualObject
函数中,相关的接口包括: get_pytorch_oneflow_res
、 get_pytorch_oneflow_tensor_res
、 oneflow_eager_run_with_graph_check
、 oneflow_tensor_eager_run_with_graph_check
、 get_oneflow_eager_res
、 get_tensor_graph_res
、 get_functional_graph_res
和 get_module_graph_test
。看一下每个接口的功能,如下表。
了解每个函数功能之后,再来看一下调用链,如下流程图所示,图中包含了 AutoTest 框架中对于 Graph 模式存在 nn.module
、nn.functional
和 tensor
三个类别的方法如何处理,对应图中的三个灰色框。
在分析了 nn.module
、nn.functional
和 tensor
三个类别的处理方法之后,其中,自动测试 Graph 时也存在一个反向的梯度测试,但是并没有取出 tensor 对应的梯度,也就是说,可以保证后向执行是正常的,没有检查 grad 值。
对于使用方法,当 @autotest()
打开 auto_backward=True
时(默认就是打开的),不仅会跑 Eager 的 Backward 测试(这里会对梯度结果做比较),还会跑对应 Graph 的 Backward 测试(这里不做梯度比较)。
对应上述描述的代码,可以在文章 3.2 部分的代码中找到:
此外,对于一些算子 inplace 版本的 Graph 检查,需要对输入做深拷贝,来保证 Graph 和 Eager 的 input 始终一致。如下代码中,get_args_copy
(在 torch_flow_dual_object.py 中)分别对普通参数和关键字参数做了 deepcopy。
类似的,在 Graph 单测中,存在 oneflow
深拷贝为 graph_train_oneflow
的行为,主要为了防止在测试一些算子时,Eager 的值被 Eager Inplace 修改,造成 Graph 的输入和 Eager 不一致导致测试出错的情况。
最后,为了保证 tensor deepcopy 的正确性,在 OneFlow 中,copy.deepcopy
会调用 tensor 的 getState
和 setState
方法,tensor 的 state 需要同时包括 data、dtype 和 device 信息,缺一不可。具体代码见:https://github.com/Oneflow-Inc/oneflow/blob/e00ba51364ff87e39edc409be395e5ed493a4ac0/python/oneflow/framework/check_point_v2.py#L159
4、Graph 的 Debug 支持
在 3.2 的代码中,可以发现存在 if verbose:
的判断,当 verbose = True
时,会输出 Graph 的 debug 信息(如算子运行 Graph 模式后的计算结果等),当然也包括 eager 下的其他需要的调试信息。
当测试出现问题时,可以通过该功能拿到错误样例,构造最小复现代码。开启方法通过环境变量控制:ONEFLOW_TEST_VERBOSE = 1
。AutoTest 框架里这个功能更多针对开发者,OneFlow 的 Graph 针对用户也提供了调试功能。
Graph 模式支持了学习率的调试输出,开启方法和 Eager 相同。
此外,调用 Graph 对象的 debug 方法,就开启了 Graph 的调试模式。
v_level=0
时,只输出最基础的警告和构图阶段信息,如构图时间。v_level=1
时,将额外打印每个nn.Module
的构图信息。v_level=2
时,在构图阶段,将额外打印每个 Op 的创建信息,包括名称、输入内容、设备和 SBP 信息等。v_level=3
时,将额外打印每个 Op 更详细的信息,如与代码位置有关的信息,方便定位代码问题。
这部分更详细的内容可以在https://docs.oneflow.org/master/basics/08_nn_graph.html#graph_3
中发现。
5、总结
AutoTest 框架的灵活性和易用性都比较强,本文主要介绍了 Graph 模式如何伴随 Eager 模式做算子对齐和 Graph 的自动测试个性化内容。Eager 到 Graph 的 Local ops 执行测试覆盖也已经在 OneFlow v0.7.0 中完成,在 0.8 版本中将会保证 Graph Global ops 单测的正确性。此外,静态图的 debug 和其他功能等将更完备。欢迎大家学习或者使用。
相关链接
https://github.com/Oneflow-Inc/oneflow
https://github.com/pytorch/pytorch
其他人都在看
OneFlow学习笔记:从Functor到OpExprInterpreter
OneFlow v0.7.0发布:全新分布式接口,LiBai、Serving等一应俱全
欢迎下载体验 OneFlow v0.7.0 最新版本:https://github.com/Oneflow-Inc/oneflow/
版权声明: 本文为 InfoQ 作者【OneFlow】的原创文章。
原文链接:【http://xie.infoq.cn/article/8d19231090e78a5b35453601a】。文章转载请联系作者。
评论