深度學習編譯器Data Flow和Control Flow

本文介紹了一下深度學習框架的Data Flow和Control Flow，基於TensorFlow解釋了TensorFlow是如何在靜態圖中實現Control Flow的。支持在Python層直接寫Control Flow的動態圖，最后基於Pytorch介紹了如何將Python層的Control Flow導出到TorchScript模型以及ONNX模型。

1. 前言

1.1. DataFlow

以TensorFlow1.x為例介紹一下DataFlow。

要實現一個的邏輯，都是一個簡單的實數，如果用Python實現非常簡單：

#coding=utf-8



import os



def cal(a, b, c):

    res = (a + b) * c

    print(res)

    return res



print(cal(1.0, 2.0, 3.0))

輸出結果是9.0。使用tf1.31.1同樣實現這個過程：

import tensorflow as tf



def cal(a, b, c):

    add_op = a + b

    print(add_op)

    mul_op = add_op * c



    init = tf.global_variables_initializer()

    sess = tf.Session()

    

    sess.run(init)

    mul_op_res = sess.run([mul_op])



    return mul_op_res



a = tf.constant(1.0)

b = tf.constant(2.0)

c = tf.constant(3.0)



print(cal(a, b, c))

同樣代碼的輸出是9.0。然后這兩個示例是為了解釋像TensorFlow這種框架，計算圖是一個計算流圖，由數據驅動的。在上面的程序中，可以發現如果打印add_op獲得的結果是一個Tensor：

Tensor("add:0", shape=(), dtype=float32

TensorFlow1.x實現的這個計算函數，先在內存中構造了一個數據流圖：

 

 

 上面tensorflow程序對應的數據流圖

Python的實現，實際上在執行res = (a + b) * c代碼時，已經計算出了res的值，因為Python這種過程語言的數學計算是由代碼驅動的。TensorFlow不一樣，先構造了數據流圖，然后對這個計算流圖進行綁定數據，讓這個數據在這個圖里面流起來，這是顯示調用sess.run獲得輸出的。

像TensorFlow這種基於數據流圖（DataFlow）進行計算的深度學習框架不少，如早期的Theano，2020年開源的國內深度學習框架OneFlow，PaddlePaddle1.x 初級版本都是基於數據流圖的。當然更多人稱為靜態圖。

1.2. Control Flow

將結合TensorFlow1.x的Control Flow解析一下Control Flow的難點，及TensorFlow的一些解決方案。這里的內容理解主要基於這篇博客(https://www.altoros.com/blog/logical-graphs-native-control-flow-operations-in-tensorflow/)，可以去查看原文。

在計算機科學中，控制流（Control Flow）定義了獨立語句，指令，函數調用等執行或者求值的順序。舉個例子，要實現一個本機控制流，即需要根據函數A的輸出值選擇運行函數B或者C中的一個：

 

 

 一個Control Flow的例子

然后要實現這個控制流，最Naive的方式在是Python端寫if/else語句，即Python端的Control Flow，然后在不同條件下使用session.run()，求取不同分支的值。對於TensorFlow是這樣：

 

 

 這里獲取A的值只是反饋回來

然后這個Python層的Control Flow不會在計算圖中被表示出來，即：

 

 

 黃色部分在計算圖中實際上是被刪掉了，因為早期的TensorFlow無法表示這種控制邏輯

可以看到上面的實現是比較爛的，這是因為使用sess.run對A進行求值后，沒做任何修改又放回了原始的計算圖，TensorFlow 計算圖與 Python 交換數據頻繁時，會嚴重拖慢運算速度。除了性能問題，在Python層做Control Flow，會發現在計算圖中，沒有表示 Python 邏輯，如果將 graph 導出，實際上是看不到這些 if/else 語句的，因此網絡結構信息會丟失。

這個問題趟過Pytorch導出ONNX的應該知道，如果想導出一個完整的檢測模型，帶了NMS后處理，必須找一張可以正常輸出目標的圖片作為輸入。如果隨機輸出，很可能后處理那部分在導出時就會丟掉，因為在Pytorch實現檢測模型時，在Python層用了if這種Control Flow。Pytorch在導出ONNX模型時，根據輸入跑一遍模型即tracing（這是以前的版本的做法，新版本的TensorFlow已經支持導出Python層的Control Flow），記錄這個過程中發生了哪些操作。如果實現模型的過程中，有Python層的Control Flow（基於tracing機制），必然有一部分節點會丟棄。

Pytorch官方文檔指出，當導出ONNX時，如果想導出Python層的控制流到計算圖中，就需要包一層@jit.script。

大概就是如果想在Pytorch里面導出含有Python層控制流的模型時導出ONNX會丟失控制流，如果需要保留建議導出TorchScript模型或者使用基於script模型的導出方式

 

 像Pytorch這種動態圖框架，可以方便的使用Python層的Control Flow，但TensorFlow在1.x時代，為了解決這個問題，花費了不少努力，即TensorFlow1.x的原生控制流。

TensorFlow的原生控制流

TensorFlow提供了幾個運算符用於原生控制流，如下：

 

 

 TensorFlow提供了幾個運算符用於原生控制流

使用這些原生控制流好處是什么呢？

高效。TensorFlow 計算圖與 Python 交換數據比較慢，計算圖如果是端到端的，才能將數據傳輸開銷降到最低，運行速度更快。

 靈活。靜態計算圖可以使用動態模塊加強，計算圖邏輯是自包含的。Pytorch目前比TensorFlow更受歡迎，主要原因就是前者為動態計算圖，可以在運行時修改計算圖。TensorFlow 利用控制流可以在一個靜態定義的計算圖中，實現類似動態計算圖的功能。

 兼容。通過 TensorBoard 調試和檢查計算圖，無縫通過 TensorFlow Serving 部署，也可以利用自動微分，隊列和流水線機制。

 

控制依賴

TensorFlow會記錄每一個運算符的依賴，然后基於依賴進行調度計算。一個運算符當且僅當依賴都完成后，才會執行一次。任何兩個完成依賴的運算符，可以以任意順序進行。但這種設定可能會引發競爭，比如：

 

 

 控制依賴引發競爭

其中 var 為一個變量，在對 bot 求值時，var 本身自增 2，將自增后的值返回。這時 top 語句執行順序就會對 out 結果產生不同影響，結果不可預知。

為了解決這個問題，開發者可以人為的加入bot和top的依賴關系，讓指定運算符先完成，如下圖所示：

 

 

 人為的加入bot和top的依賴關系，讓指定運算符先完成

如果需要保證讀取的值最新，需要新增下圖中虛線箭頭表示的依賴關系，即下圖中上方藍色圓圈依賴下方藍色圓圈的運算完成，才能進行計算。

 

 

 加入依賴關系后，計算圖長這樣

條件分支

接下來看條件分支，即TensorFlow如何處理在這一節開頭提出來的那個例子？

 

 

 TensorFlow提供了兩個條件控制OP，即tf.cond和tf.case

下面的代碼中，利用了tf.cond實現條件分支，在 a < b 為真，對 out 求值會執行 tf.add(3, 3)；否則，執行 tf.square(3)。

 

 

 使用tf.cond實現條件分支

上面這段代碼等價於：tf.cond(a < b, lambda: tf.add(3, 3), lambda: tf.sqaure(3))

然后生成的計算圖如下所示：

 

 

 帶有條件控制流的計算圖

當並列的分支比較多時，可以使用tf.case來處理，例如：

 

 

 並列的條件分支>2個時，使用tf.case來控制

循環

TensorFlow提供了tf.while_loop來構造循環塊，感覺和RNN類似的結構有這個需求，例如：

 

 

 tf.while_loop可以實現循環控制流解決RNN這種計算圖結構的控制邏輯

下面的代碼實現了一個基礎的循環例子，即循環100次。

 

 

 使用tf.while_loop在靜態圖中實現循環控制流

總的來說，TensorFlow應該是首個將Control Flow引入到計算圖中的深度學習框架，不是像動態圖框架那樣直接在Python層去做Control Flow，這方面必須給予一定的尊重。即使Pytorch目前在學術界已經比TensorFlow更加流行，但基於TensorFlow演化的各種工業級項目仍然發揮着作用。

3. Pytorch中的Control Flow

在Pytorch這種動態圖框架中，支持直接在Python端寫Control Flow，並且可以將這些控制邏輯放到計算圖中。這里以TorchScript為例，當嘗試將Pytorch模型轉為TorchScript時，有兩種方式，一種是trace，另外一種是script。對於trace模式，適合Python層沒有Control Flow的計算圖，舉例如下：

#coding=utf-8

import torch

import torch.nn as nn



class MyModule(nn.Module):

    def __init__(self):

       super(MyModule,self).__init__()

       self.conv1 = nn.Conv2d(1,3,3)

    def forward(self,x):

       x = self.conv1(x)

       return x



model = MyModule()  # 實例化模型

trace_module = torch.jit.trace(model,torch.rand(1,1,224,224)) 

print(trace_module.code)  # 查看模型結構

output = trace_module (torch.ones(1, 1, 224, 224)) # 測試

print(output)

# trace_modult('model.pt') 

打印trace_module的代碼可以看到：

def forward(self,

    input: Tensor) -> Tensor:

  return (self.conv1).forward(input, )

而script模式則適用於計算圖在Python層有Control Flow的情況，比如：

#coding=utf-8

import torch

import torch.nn as nn



class MyModule(nn.Module):

    def __init__(self):

        super(MyModule,self).__init__()

        self.conv1 = nn.Conv2d(1,3,3)

        self.conv2 = nn.Conv2d(2,3,3)



    def forward(self,x):

        b,c,h,w = x.shape

        if c ==1:

            x = self.conv1(x)

        else:

            x = self.conv2(x)

        return x



model = MyModule()



# 這樣寫會報錯，因為有控制流

# trace_module = torch.jit.trace(model,torch.rand(1,1,224,224)) 



# 此時應該用script方法

script_module = torch.jit.script(model) 

print(script_module.code)

output = script_module(torch.rand(1,1,224,224))

打印script_module的代碼可以看到TorchScript模型包含了在上面Python層定義的Control Flow：

def forward(self,

    x: Tensor) -> Tensor:

  b, c, h, w, = torch.size(x)

  if torch.eq(c, 1):

    x0 = (self.conv1).forward(x, )

  else:

    x0 = (self.conv2).forward(x, )

  return x0

然后來實驗一下將上面帶有Control Flow的Module導出ONNX，這里以Pytorch官方文檔提供的一個帶循環的Control Flow的示例為例：

import torch



# Trace-based only



class LoopModel(torch.nn.Module):

    def forward(self, x, y):

        for i in range(y):

            x = x + i

        return x



model = LoopModel()

dummy_input = torch.ones(2, 3, dtype=torch.long)

loop_count = torch.tensor(5, dtype=torch.long)



torch.onnx.export(model, (dummy_input, loop_count), 'loop.onnx', verbose=True)

這樣就可以成功導出名字為loop的ONNX模型，使用Netron可視化軟件打開看一下：

 

 

 

可以看到直接導出Module，Python層的控制邏輯被丟掉（即for循環被完全展開），這是因為Pytorch在導出ONNX的時候默認使用了tracing機制

而當使用script模式時，導出的ONNX就會保留Python層的Control Flow並將其轉換成ONNX中的Loop OP。示例代碼以及Netron可視化結果如下：

import torch

# Mixing tracing and scripting



@torch.jit.script

def loop(x, y):

    for i in range(int(y)):

        x = x + i

    return x



class LoopModel2(torch.nn.Module):

    def forward(self, x, y):

        return loop(x, y)



model = LoopModel2()

dummy_input = torch.ones(2, 3, dtype=torch.long)

loop_count = torch.tensor(5, dtype=torch.long)

torch.onnx.export(model, (dummy_input, loop_count), 'loop.onnx', verbose=True,

                  input_names=['input_data', 'loop_range'])

Pytorch模型中在Python層定義的Control Flow被保留下來了

4. 總結

這篇文章介紹了一下深度學習中的Data Flow和Control Flow，然后介紹了一下將Pytorch模型轉為TorchScript的兩種模式，並探索了要將Pytorch的Python層的Control Flow轉換為ONNX應該怎么做。

5. 參考文獻

 

https://mp.weixin.qq.com/s/Kt4xDLo-NRui8Whl0DqcSA

 

 

https://blog.csdn.net/lvxingzhe123456/article/details/82597095

 

 

https://www.altoros.com/blog/logical-graphs-native-control-flow-operations-in-tensorflow/

 

 

https://mp.weixin.qq.com/s/6uVeEHcQeaPN_qEhHvcEoA