深度學習框架如何優雅地做算子對齊任務?

語言: CN / TW / HK

直播提前約!!!

今晚19:30 ,OneFlow工程師教你 用AI算法拯救渣畫質照片 ,歡迎參與。

撰文 | BBuf

之前回答過 「如何為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

1

傳統的算子對齊方式

不侷限於OneFlow,任何組織或者個人編寫的深度學習訓練框架都需要驗證算子的實現正確性。那麼,深度學習框架中驗證算子正確性的一般做法是什麼呢?

以百度的PaddlePaddle為例,在驗證算子正確性時一般是根據調用其它標準庫獲得的結果(比如卷積算子的驗證就調用cudnn的卷積,erf算子的驗證就調用了scipy的erf)或者直接使用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文件, 並且這個AutoTest框架可以輕易移植到其它任何深度學習框架去做算子對齊任務。

1.https://github.com/Oneflow-Inc/oneflow/blob/v0.6.0/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py

2.https://github.com/Oneflow-Inc/oneflow/blob/v0.6.0/python/oneflow/test_utils/automated_test_util/generators.py

2

算子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算子實現是否存在問題。

3

算子AutoTest框架實現思路

瞭解了AutoTest框架的使用方法,這裏來講解一下AutoTest框架的實現思路。從上面的用法可以大概可以猜到AutoTest框架在實現時會分成兩部分,一部分是如何產生隨機數據,另外一部分是用AutoTest部分的程序並記錄和比較中間tensor以及對應的梯度tensor的形狀和數值。

3.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 類),也可以忽略這個 annotation 直接在 _calc_value 構造相應類型的隨機結果(如 random_device 類)。

3.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:\n{self.pytorch}\n\nOneFlow object:\n{self.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  接口時傳入的是 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)         # print(oneflow_attr)         new_name = f"{self.name}.{key}"         return GetDualObject(new_name, pytorch_attr, oneflow_attr) # flatten的AutoTest程序 @autotest(auto_backward=False) def test_against_pytorch(test_case):     m = torch.nn.Flatten(         start_dim=random(1, 6) | nothing(), end_dim=random(1, 6) | nothing()     )     m.train(random())     device = random_device()     m.to(device)     x = random_pytorch_tensor().to(device)     y = m(x)     return y 

然後看一下 __getattr__ 中key的打印結果:

nn Flatten train to to

可以看到被 autotest() 裝飾器修飾的測試程序中的PyTorch或者OneFlow的 nn.Module 或者其它函數都重寫了這個方法,它將這些nn.Module或者其它函數的參數和屬性都取出來並同樣使用 GetDualObject 返回一個新的 DualObject 對象,我們可以打印一下 Flatten 這個 nn.Module 對應的 DualObject 對象是什麼:

PyTorch object: <bound method Module.train of Flatten(start_dim=1, end_dim=-1)> OneFlow object: <bound method Module.train of Flatten(start_dim=1, end_dim=-1)> 

GetDualObject 這個函數就是根據傳入的 Pytorch 以及 OneFlow 對象和它們的名字來生成一個 DualObject 對象。 GetDualObject 這個函數會為high level的PyTorch重寫傳入的原始PyTorch以及OneFlow對象的 __call__ 魔法函數,最後返回一個 DualObject 對象,這個過程還包含了跳過一些不需要關注的魔法函數以及檢查傳入對象的屬性是否合法和基於nn.Module和其它API默認參數的類型對 generator 繼承類產生的隨機數據綁定特定類型的工作( get_args 函數中完成)。這裏還有一句對於Tensor方法的特判,因為Tensor方法的調用方式(通過 getattr )和其它Module和函數不同(通過 __call__ )。

GetDualObject 的實現思路大致就是這樣,代碼比較長這裏就不貼了,感興趣可以在這裏查看:

https://github.com/Oneflow-Inc/oneflow/blob/v0.6.0/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py#L195-L401  

最後,我們看一下 autotest() 裝飾器的實現:

def autotest(     n=20,     auto_backward=True,     rtol=0.0001,     atol=1e-05,     check_graph=True,     check_allclose=True, ):     verbose = os.getenv("ONEFLOW_TEST_VERBOSE") is not None     def deco(f):         @functools.wraps(f)         def new_f(test_case):             nonlocal n             loop_limit = n * 20             loop = 0             while n > 0:                 clear_note_fake_program()                 if loop > loop_limit:                     raise ValueError("autotest stuck in an endless loop!")                 dual_modules_to_test.clear()                 dual_objects_to_test.clear()                 try:                     global testing                     testing = True                     global testing_graph                     if check_graph:                         testing_graph = True                     res = f(test_case)                     testing = False                     testing_graph = False                 except (PyTorchDoesNotSupportError, BothDoNotSupportError) as e:                     if verbose:                         print(f"{f.__name__}")                         print(e)                     loop += 1                     continue                 if res is not None:                     if not isinstance(res, collections.abc.Sequence):                         res = [res]                     func_outputs = res                     for x in res:                         if auto_backward:                             if isinstance(x.pytorch, torch_original.Tensor):                                 call_tensor_id.append(id(x.pytorch))                                 x.sum().backward()                         dual_objects_to_test.append(x)                 for x in dual_modules_to_test:                     for key in x.pytorch.state_dict().keys():                         if key not in x.oneflow.state_dict().keys():                             warnings.warn(f"oneflow module don't have `{key}`")                             continue                         vis_parameters[key] = x.pytorch.state_dict()[key]                         dual_objects_to_test.append(                             GetDualObject(                                 "unused",                                 getattr(x.pytorch, key),                                 getattr(x.oneflow, key),                             )                         )                         call_tensor_id.append(id(getattr(x.pytorch, key)))                         dual_objects_to_test.append(                             GetDualObject(                                 "unused",                                 getattr(x.pytorch, key).grad,                                 getattr(x.oneflow, key).grad,                             )                         )                         call_tensor_id.append(id(getattr(x.pytorch, key).grad))                 for x in dual_objects_to_test:                     if (                         isinstance(x.pytorch, torch_original.Tensor)                         and id(x.pytorch) not in call_tensor_id                     ):                         vis_tensor.append(x.pytorch)                 # check eager                 for x in dual_objects_to_test:                     if check_allclose:                         test_case.assertTrue(check_equality(x, rtol=rtol, atol=atol), x)                     if verbose:                         print(f"{f.__name__} test eager passed.")                                     n -= 1                 loop += 1         return new_f     return deco 

這個裝飾器的 res = f(test_case) 這行代碼會執行這個裝飾器修飾的自動測試程序,會在給定輸入的情況下去分別運行PyTorch和OneFlow的程序獲得所有中間的輸出tensor,包括tensor的梯度,並將它們記錄到 dual_modules_to_test 這個列表。再遍歷這個列表裏面的每個tensor,比較數值和shape是否完全一樣。比較函數實現在:

https://github.com/Oneflow-Inc/oneflow/blob/v0.6.0/python/oneflow/test_utils/automated_test_util/torch_flow_dual_object.py#L565-L599 

原理就是拿到tensor的numpy數據進行比較。 autotest()  裝飾器還有幾個參數可以調整,可以控制測試是否執行反向,執行次數,以及最後結果對比的精度閾值。

4

自動生成出BUG的程序和數據

上面介紹完了AutoTest框架的原理和使用方法,這裏再展示一下基於AutoTest框架如何拿到可復現BUG的程序以及對應的輸入tensor和參數等。原理很簡單,就是把 GetDualObject 過程中使用的api記錄下來拼起來就構成一個完整的程序,這裏展示一下在CI中的效果。

https://github.com/Oneflow-Inc/oneflow/runs/4760189461?check_suite_focus=true 

這個例子展示了在某次CI過程中,OneFlow的 conv_transpose2d 算子和PyTorch的 conv_transpose2d 算子在某個case下沒有對齊,那麼CI在報吿這個錯誤時也輸出了對應的復現代碼和數據,可以方便框架開發者進行定位和判斷:

自動測試框架在算子和PyTorch沒對齊時會輸出復現程序和數據

除此之外,這個AutoTest框架目前不僅負責Eager算子的測試,還被我們擴展到支持nn.Graph和Eager Consistent等多種情況,極大的方便了框架開發者。

5

總結

這篇文章介紹了OneFlow的算子AutoTest框架,提供了一個深度學習優雅地做算子對齊的方法,使得開發者和用户可以像寫PyTorch那樣方便寫測試程序。AutoTest框架的靈活性和易用性都比較強,歡迎大家學習或者使用。

原文首發於:https://zhuanlan.zhihu.com/p/458111952

其他人都在看

點擊“ 閲讀原文 ,歡迎下載體驗OneFlow新一代開源深度學習框架