深度学习框架如何优雅的做算子对齐任务？
Posted just_sort
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了深度学习框架如何优雅的做算子对齐任务？相关的知识，希望对你有一定的参考价值。
0x0. 前言
之前回答过如何为PyTorch做贡献的知乎问题，原贴见：https://www.zhihu.com/question/502301777/answer/2248950419 。回答提到了去年在OneFlow开发一些算子时，基于算子AutoTest框架找到了一些PyTorch算子的bug，并给PyTorch做出了反馈或修复。但这个回答没有介绍这个AutoTest框架长什么样子，以及它背后的原理。因此，这篇文章就用来介绍OneFlow的算子AutoTest框架看一下OneFlow深度学习框架在算子开发过程中是如何优雅的做算子对齐任务的（由@大缺弦 开发，后经我和其它同事进行扩展和丰富功能形成今天的形态）。这个AutoTest框架也可以很轻易移植到其它深度学习训练框架使用，代码实现在https://github.com/Oneflow-Inc/oneflow/blob/v0.6.0/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py。
0x1. 传统的算子对齐方式
不局限于OneFlow，任何组织或者个人编写的深度学习训练框架都需要验证算子的实现正确性。那么，深度学习框架中验证算子正确性的一般做法是什么呢？以百度的PaddlePaddle为例，在验证算子正确性时一般是根据调用其它标准库获得结果（比如卷积算子的验证就调用cudnn的卷积，erf算子的验证就调用了scipy的实现）或者直接使用numpy模拟的计算结果来对不验证（比如full算子的验证即为numpy模拟）。在PyTorch的测试样例中还有硬编码一些测试样例的方式，也即将固定输入样例的标准答案和算子输出的结果进行对比，以此判断算子的正确性。
这些方法都没有什么问题，但在编写测试时需要不少的人力并且在算子开发初期可能有一些corner case会容易想不到。以OneFlow为例，由于算子的行为是对齐PyTorch，如果要验证转置卷积Op在各种情况下的正确性，那么什么样的测试代码才可以全面验证呢？一种做法是将每个参数都枚举出来：
import torch
import numpy as np
import oneflow as flow
for N in range(1, 5):
for C_in in range(1, 10):
for L_in in range(1, 10):
for H_in in range(1, 10):
for C_out in range(1, 10):
for Ksize in range(1, 10):
for Pad in range(1, 10):
for Dilation in range(1, 10):
for Stride in range(1, min(L_in, H_in)):
for OutPad in range(1, min(Dilation, Stride)):
try:
torch_input = torch.randn(N, C_in, L_in, H_in)
flow_input = flow.tensor(torch_input.numpy())
torch_input.requires_grad = True
flow_input.requires_grad = True
torch_m = torch.nn.ConvTranspose2d(in_channels=C_in, out_channels=C_out, kernel_size=Ksize, padding=Pad, stride=Stride,
output_padding=(OutPad), dilation=Dilation, bias=False)
flow_m = flow.nn.ConvTranspose2d(in_channels=C_in, out_channels=C_out, kernel_size=Ksize, padding=Pad, stride=Stride,
output_padding=(OutPad), dilation=Dilation, bias=False)
flow_m.weight.data = flow.tensor(torch_m.weight.data.detach().numpy(), requires_grad=True)
torch_out = torch_m(torch_input)
flow_out = flow_m(flow_input)
torch_out = torch_out.sum()
flow_out = flow_out.sum()
assert(np.allclose(torch_out.detach().numpy(), flow_out.detach().numpy(), 1e-06, 1e-06)), "forward not equal"
torch_out.backward()
flow_out.backward()
print(torch_input.grad.detach().numpy())
print(flow_input.grad.detach()[:N, :C_in, :L_in, :H_in].numpy())
assert(np.allclose(torch_input.grad.detach().numpy(), flow_input.grad.detach()[:N, :C_in, :L_in, :H_in].numpy(), 1e-03, 1e-03)), "backward not equal"
except Exception as e:
print('Input Param Error')
但这种做法虽然验证得比较全面但同样有缺点。首先枚举的上界如何确定？如果给了一个大的上界，那么这个算子的验证时间会非常长，不利于在CI流程中使用。如果上界很小就可能忽略一些corner case，导致测试仍然不会全面并增加算子出bug的风险。
基于算子测试的这些问题，同事 @大缺弦 开发了一个算子AutoTest框架，用于解决OneFlow算子和PyTorch算子对齐的问题。后来我在此基础上又为这个AutoTest框架丰富了其它的一些功能，感觉目前已经比较好使，接下里做一个全面介绍。
整个AutoTest框架只有2个Python文件，即：https://github.com/Oneflow-Inc/oneflow/blob/v0.6.0/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py 和 https://github.com/Oneflow-Inc/oneflow/blob/v0.6.0/python/oneflow/test_utils/automated_test_util/generators.py 。并且这个AutoTest框架可以轻易移植到其它任何深度学习框架去做算子对齐任务。
0x2. 算子AutoTest框架用法
在介绍原理之前，我们先看一下AutoTest框架的用法。以上面的反卷积算子为例，使用了AutoTest框架之后就可以用下面的代码来完成算子对齐测试：
@autotest()
def test_deconv2d_with_random_data(test_case):
channels = random(1, 6)
m = torch.nn.ConvTranspose2d(
in_channels=channels,
out_channels=random(1, 20),
kernel_size=random(1, 4),
stride=random() | nothing(),
padding=random(1, 3).to(int) | nothing(),
dilation=random(1, 5) | nothing(),
groups=random(1, 5) | nothing(),
padding_mode=constant("zeros") | nothing(),
)
m.train(random())
device = random_device()
m.to(device)
x = random_pytorch_tensor(ndim=4, dim1=channels).to(device)
y = m(x)
return y
熟悉PyTorch的小伙伴可以发现这个算子测试代码和PyTorch的代码风格基本一样。的确，AutoTest框架相当于是一个high level的PyTorch，它的接口和PyTorch一样，但对于给定的输入会分别用OneFlow和PyTorch运行一遍，记录运行过程中得到的每个tensor以及对应梯度tensor的值，再对这些OneFlow和PyTorch分别产生的tensor检查一遍数值形状是否完全相同，以完成自动测试工作，我们后面会细讲。
我们可以再看一个测试matmul算子的例子：
@autotest()
def test_flow_matmul_with_random_data(test_case):
k = random(1, 6)
x = random_pytorch_tensor(ndim=2, dim1=k)
y = random_pytorch_tensor(ndim=2, dim0=k)
z = torch.matmul(x, y)
return z
我们基于random_pytorch_tensor方法构造了两个随机tensor x和y，它们的维度分别是[m, k]和[k, n]，这些维度的值都是随机生成的。
执行上述两个测试例子，自动测试框架会自动帮我们随机出各种合法参数组合成的Op，并基于数值和类型完全相同的输入Tensor（PyTorch和OneFlow各有一份）分别运行PyTorch和OneFlow的代码，并完成算子的自动测试。由于自动测试框架的用法对齐了PyTorch用法，我们在开发算子之后编写测试样例将非常简单。不用再引入其它的标准库或者使用Numpy去模拟一遍算子的前向反向计算过程等，解放了生产力。
并且测试的时候只要次数足够多，就可以很大概率的覆盖到一些OneFlow算子和PyTorch算子无法对齐的样例，这个时候如果能拿到对应的复现样例就可以帮助我们确定OneFlow算子实现是否存在问题。
0x3. 算子AutoTest框架实现思路
了解了AutoTest框架的使用方法之后，这里来讲解一下AutoTest框架的实现思路。从上面的用法可以大概可以猜到AutoTest框架在实现时会分成两部分，一部分是如何产生随机数据，另外一部分是运AutoTest部分的程序并记录和比较中间tensor以及对应的梯度tensor的形状和数值。
0x3.1 如何产生随机数据？
这里说的随机数据不仅指的是随机的输入tensor，还包含Op的属性参数比如上面反卷积Op测试例子中的kernel_size=random(1, 4)就实现了指定kernel_size将会在[1, 4)这个区间进行取值。
这部分实现在https://github.com/Oneflow-Inc/oneflow/blob/v0.6.0/python/oneflow/test_utils/automated_test_util/generators.py 这个文件里。首先我们看一下这个文件导出了哪些接口：
__all__ = [
"random_tensor",
"random_bool",
"random_device",
"random",
"random_or_nothing",
"oneof",
"constant",
"nothing"
]
这些接口都是继承了generator基类用来产生随机数据结构的类，这里的数据结构既可以是内置类型如int，也可以是自定义数据类型比如tensor。AutoTest框架所有的参数的随机性都是基于这些方法来做到的，我们看一下generator基类的实现：
class generator:
def __init__(self, children):
self.children = children
self._value = None
def _init(self):
self._value = None
for x in self.children:
x._init()
def eval(self):
self._init()
return self.value()
def _calc_value(self):
raise NotImplementedError()
def value(self):
if self._value is None:
self._value = self._calc_value()
return self._value
def size(self):
return 1
def __or__(self, other):
other = pack(other)
return oneof(
self, other, possibility=self.size() / (self.size() + other.size())
)
def __ror__(self, other):
return self | other
def __add__(self, other):
return add(self, other)
def __radd__(self, other):
return self + other
def __sub__(self, other):
return self + neg(other)
def __rsub__(self, other):
return neg(self - other)
def __mul__(self, other):
return mul(self, other)
def __rmul__(self, other):
return self * other
def to(self, annotation):
self._to(annotation)
for x in self.children:
x.to(annotation)
return self
def _to(self, annotation):
pass
这个类不仅持有了_calc_value，value，eval等和取值有关的函数，还持有size这个反应生成数据个数的函数。另外还持有了一系列的魔法函数，让不同的generator子类可以互相组合，提升了自动测试框架书写的灵活性。最后还有一个to成员函数，这个函数被继承generator基类的类重写，用来确定这个随机数据结构的数值类型。
所有的generator派生类都继承了generator基类，并重写其中的__init__，__calc_value，size，_to等成员函数。比如nothing这个generator就是直接重写_calc_value函数，并在其中返回一个什么都不做的类的实体。
class Nothing:
pass
class nothing(generator):
def __init__(self):
super().__init__([])
def _calc_value(self):
return Nothing()
再例如，random这个generator的派生类的定义如下：
class random(generator):
def __init__(self, low=1, high=6):
self.low = pack(low)
self.high = pack(high)
super().__init__([self.low, self.high])
self.annotation = None
def _to(self, annotation):
if self.annotation is not None:
return
if hasattr(annotation, "__origin__"):
# PyTorch _size_2_t and similar types are defined by type variables,
# leading to unexpected __args__ and __origin__
#
# >>> _size_2_t = Union[T, Tuple[T, T]][int]
# >>> _size_2_t.__origin__
# typing.Union[~T, typing.Tuple[~T, ~T]]
#
# So recreate a new annotation object by repr and eval
#
# >>> _size_2_t
# typing.Union[int, typing.Tuple[int, int]]
# >>> _size_2_t_new = eval(repr(annotation))
# >>> _size_2_t_new.__origin__
# typing.Union
annotation = eval(repr(annotation))
self.annotation = annotation
def _generate(self, annotation):
if hasattr(annotation, "__origin__"):
if annotation.__origin__ is Union:
x = random_util.choice(annotation.__args__)
return self._generate(x)
if annotation.__origin__ is Tuple or annotation.__origin__ is py_tuple:
return [self._generate(x) for x in annotation.__args__]
else:
raise NotImplementedError(
f"Not implemented annotation annotation in random, type(annotation.__origin__) is type(annotation.__origin__)"
)
low, high = self.low.value(), self.high.value()
if annotation == int:
val = int(rng.integers(low, high))
elif annotation == float:
val = float(rng.random() * (high - low) + low)
elif annotation == bool:
val = random_util.choice([True, False])
else:
raise NotImplementedError(
f"Not implemented annotation annotation in random"
)
return val
def _calc_value(self):
return self._generate(self.annotation)
def random_or_nothing(low, high):
return oneof(random(low, high), nothing(), possibility=2 / 3)
这里需要注意的一点是，持有annotation属性的generator派生类的是可以通过to来更新annotation属性（如random类），也可以忽略这个参数直接在_calc_value构造相应类型的随机结果（如random_device类）。
0x3.2 AutoTest核心实现
AutoTest框架的核心实现在https://github.com/Oneflow-Inc/oneflow/blob/v0.6.0/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py 这个文件。这个文件最后2行代码是：
torch = GetDualObject("", torch_original, flow)
__all__ = ["autotest", "random_pytorch_tensor"]
这行代码torch = GetDualObject("", torch_original, flow) 里面的torch_original表示原始的PyTorch框架，而使用GetDualObject获得的torch表示是对原始的PyTorch和OneFlow进行了一个封装，变成了一个high level的PyTorch。因此，这里最关键的实现就是GetDualObject这个函数，我们先不关注这个函数具体在做什么，而是它返回了什么。查看代码可以发现这个函数返回了一个DualObject类对象，我们先研究一下这个类：
class DualObject:
def __init__(self, name, pytorch, oneflow):
self.name = name
self.pytorch = pytorch
self.oneflow = oneflow
if isinstance(pytorch, torch_original.nn.Module):
state_dict = pytorch.state_dict()
state_dict = k: v.detach().cpu().numpy() for (k, v) in state_dict.items()
oneflow.load_state_dict(state_dict, strict=False)
if testing:
dual_modules_to_test.append(self)
if isinstance(pytorch, torch_original.Tensor):
if testing:
dual_objects_to_test.append(self)
def __repr__(self):
return f"PyTorch object:\\nself.pytorch\\n\\nOneFlow object:\\nself.oneflow"
def __getattr__(self, key):
pytorch_attr = getattr(self.pytorch, key)
oneflow_attr = getattr(self.oneflow, key)
new_name = f"self.name.key"
global call_pytorch
call_pytorch = self.pytorch
return GetDualObject(new_name, pytorch_attr, oneflow_attr)
在__init__中传入了类对象名和pytorch/oneflow两个对象，在导出high level的PyTorch的时候传入的是torch_original和flow，而在导出random_pytorch_tensor API时传入的是pytorch_tensor和oneflow_tensor。这里不妨先看一下random_pytorch_tensor这个函数的实现：
def random_pytorch_tensor(
ndim=None,
dim0=1,
dim1=None,
dim2=None,
dim3=None,
dim4=None,
low=0,
high=1,
dtype=float,
requires_grad=True,
):
if isinstance(requires_grad, generator):
requires_grad = requires_grad.value()
pytorch_tensor = (
random_tensor(ndim, dim0, dim1, dim2, dim3, dim4, low, high, dtype)
.value()
.requires_grad_(requires_grad and dtype != int)
)
flow_tensor = flow.tensor(
pytorch_tensor.detach().cpu().numpy(),
requires_grad=(requires_grad and dtype != int),
)
return GetDualObject("unused", pytorch_tensor, flow_tensor)
可以看到它和导出high level PyTorch的实现一样，也是通过调用GetDualObject来获得了一个对象。再回到DualObject类的实现，可以发现这里分别使用了dual_modules_to_test和dual_objects_to_test这两个list来分别记录OneFlow和PyTorch的nn.Module和tensor对象。另外DualObject类还重写了__getattr__这个魔法方法，这里以Flatten为例来看看这个魔法方法获取了AutoTest程序中的那些属性：
def __getattr__(self, key):
pytorch_attr = getattr(self.pytorch, key)
oneflow_attr = getattr(self.oneflow, key)
print(key)
# print(pytorch_attr)
#以上是关于深度学习框架如何优雅的做算子对齐任务？的主要内容，如果未能解决你的问题，请参考以下文章
清华大学发布基于元算子和动态编译的深度学习框架-Jittor
业内热点清华大学发布基于元算子和动态编译的深度学习框架- Jittor