深度學(xué)習(xí)編譯器Data Flow和Control Flow
本文介紹了一下深度學(xué)習(xí)框架的Data Flow和Control Flow，基于TensorFlow解釋了TensorFlow是如何在靜態(tài)圖中實(shí)現(xiàn)Control Flow的。支持在Python層直接寫Control Flow的動(dòng)態(tài)圖，最后基于Pytorch介紹了如何將Python層的Control Flow導(dǎo)出到TorchScript模型以及ONNX模型。

1. 前言
   1.1. DataFlow
   以TensorFlow1.x為例介紹一下DataFlow。
   要實(shí)現(xiàn)一個(gè)的邏輯，都是一個(gè)簡單的實(shí)數(shù)，如果用Python實(shí)現(xiàn)非常簡單：
   #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))
輸出結(jié)果是9.0。使用tf1.31.1同樣實(shí)現(xiàn)這個(gè)過程：
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。然后這兩個(gè)示例是為了解釋像TensorFlow這種框架，計(jì)算圖是一個(gè)計(jì)算流圖，由數(shù)據(jù)驅(qū)動(dòng)的。在上面的程序中，可以發(fā)現(xiàn)如果打印add_op獲得的結(jié)果是一個(gè)Tensor：
Tensor(“add:0”, shape=(), dtype=float32
TensorFlow1.x實(shí)現(xiàn)的這個(gè)計(jì)算函數(shù)，先在內(nèi)存中構(gòu)造了一個(gè)數(shù)據(jù)流圖：

上面tensorflow程序?qū)?yīng)的數(shù)據(jù)流圖
Python的實(shí)現(xiàn)，實(shí)際上在執(zhí)行res = (a + b) * c代碼時(shí)，已經(jīng)計(jì)算出了res的值，因?yàn)镻ython這種過程語言的數(shù)學(xué)計(jì)算是由代碼驅(qū)動(dòng)的。TensorFlow不一樣，先構(gòu)造了數(shù)據(jù)流圖，然后對(duì)這個(gè)計(jì)算流圖進(jìn)行綁定數(shù)據(jù)，讓這個(gè)數(shù)據(jù)在這個(gè)圖里面流起來，這是顯示調(diào)用sess.run獲得輸出的。
像TensorFlow這種基于數(shù)據(jù)流圖（DataFlow）進(jìn)行計(jì)算的深度學(xué)習(xí)框架不少，如早期的Theano，2020年開源的國內(nèi)深度學(xué)習(xí)框架OneFlow，PaddlePaddle1.x 初級(jí)版本都是基于數(shù)據(jù)流圖的。當(dāng)然更多人稱為靜態(tài)圖。
1.2. Control Flow
將結(jié)合TensorFlow1.x的Control Flow解析一下Control Flow的難點(diǎn)，及TensorFlow的一些解決方案。這里的內(nèi)容理解主要基于這篇博客(https://www.altoros.com/blog/logical-graphs-native-control-flow-operations-in-tensorflow/)，可以去查看原文。
在計(jì)算機(jī)科學(xué)中，控制流（Control Flow）定義了獨(dú)立語句，指令，函數(shù)調(diào)用等執(zhí)行或者求值的順序。舉個(gè)例子，要實(shí)現(xiàn)一個(gè)本機(jī)控制流，即需要根據(jù)函數(shù)A的輸出值選擇運(yùn)行函數(shù)B或者C中的一個(gè)：

一個(gè)Control Flow的例子
然后要實(shí)現(xiàn)這個(gè)控制流，最Naive的方式在是Python端寫if/else語句，即Python端的Control Flow，然后在不同條件下使用session.run()，求取不同分支的值。對(duì)于TensorFlow是這樣：

這里獲取A的值只是反饋回來
然后這個(gè)Python層的Control Flow不會(huì)在計(jì)算圖中被表示出來，即：

黃色部分在計(jì)算圖中實(shí)際上是被刪掉了，因?yàn)樵缙诘腡ensorFlow無法表示這種控制邏輯
可以看到上面的實(shí)現(xiàn)是比較爛的，這是因?yàn)槭褂胹ess.run對(duì)A進(jìn)行求值后，沒做任何修改又放回了原始的計(jì)算圖，TensorFlow 計(jì)算圖與 Python 交換數(shù)據(jù)頻繁時(shí)，會(huì)嚴(yán)重拖慢運(yùn)算速度。除了性能問題，在Python層做Control Flow，會(huì)發(fā)現(xiàn)在計(jì)算圖中，沒有表示 Python 邏輯，如果將 graph 導(dǎo)出，實(shí)際上是看不到這些 if/else 語句的，因此網(wǎng)絡(luò)結(jié)構(gòu)信息會(huì)丟失。
這個(gè)問題趟過Pytorch導(dǎo)出ONNX的應(yīng)該知道，如果想導(dǎo)出一個(gè)完整的檢測模型，帶了NMS后處理，必須找一張可以正常輸出目標(biāo)的圖片作為輸入。如果隨機(jī)輸出，很可能后處理那部分在導(dǎo)出時(shí)就會(huì)丟掉，因?yàn)樵赑ytorch實(shí)現(xiàn)檢測模型時(shí)，在Python層用了if這種Control Flow。Pytorch在導(dǎo)出ONNX模型時(shí)，根據(jù)輸入跑一遍模型即tracing（這是以前的版本的做法，新版本的TensorFlow已經(jīng)支持導(dǎo)出Python層的Control Flow），記錄這個(gè)過程中發(fā)生了哪些操作。如果實(shí)現(xiàn)模型的過程中，有Python層的Control Flow（基于tracing機(jī)制），必然有一部分節(jié)點(diǎn)會(huì)丟棄。
Pytorch官方文檔指出，當(dāng)導(dǎo)出ONNX時(shí)，如果想導(dǎo)出Python層的控制流到計(jì)算圖中，就需要包一層@jit.script。
大概就是如果想在Pytorch里面導(dǎo)出含有Python層控制流的模型時(shí)導(dǎo)出ONNX會(huì)丟失控制流，如果需要保留建議導(dǎo)出TorchScript模型或者使用基于script模型的導(dǎo)出方式，
像Pytorch這種動(dòng)態(tài)圖框架，可以方便的使用Python層的Control Flow，但TensorFlow在1.x時(shí)代，為了解決這個(gè)問題，花費(fèi)了不少努力，即TensorFlow1.x的原生控制流。
TensorFlow的原生控制流
TensorFlow提供了幾個(gè)運(yùn)算符用于原生控制流，如下：

TensorFlow提供了幾個(gè)運(yùn)算符用于原生控制流
使用這些原生控制流好處是什么呢？
? 高效。TensorFlow 計(jì)算圖與 Python 交換數(shù)據(jù)比較慢，計(jì)算圖如果是端到端的，才能將數(shù)據(jù)傳輸開銷降到最低，運(yùn)行速度更快。
? 靈活。靜態(tài)計(jì)算圖可以使用動(dòng)態(tài)模塊加強(qiáng)，計(jì)算圖邏輯是自包含的。Pytorch目前比TensorFlow更受歡迎，主要原因就是前者為動(dòng)態(tài)計(jì)算圖，可以在運(yùn)行時(shí)修改計(jì)算圖。TensorFlow 利用控制流可以在一個(gè)靜態(tài)定義的計(jì)算圖中，實(shí)現(xiàn)類似動(dòng)態(tài)計(jì)算圖的功能。
? 兼容。通過 TensorBoard 調(diào)試和檢查計(jì)算圖，無縫通過 TensorFlow Serving 部署，也可以利用自動(dòng)微分，隊(duì)列和流水線機(jī)制。
控制依賴
TensorFlow會(huì)記錄每一個(gè)運(yùn)算符的依賴，然后基于依賴進(jìn)行調(diào)度計(jì)算。一個(gè)運(yùn)算符當(dāng)且僅當(dāng)依賴都完成后，才會(huì)執(zhí)行一次。任何兩個(gè)完成依賴的運(yùn)算符，可以以任意順序進(jìn)行。但這種設(shè)定可能會(huì)引發(fā)競爭，比如：

控制依賴引發(fā)競爭
其中 var 為一個(gè)變量，在對(duì) bot 求值時(shí)，var 本身自增 2，將自增后的值返回。這時(shí) top 語句執(zhí)行順序就會(huì)對(duì) out 結(jié)果產(chǎn)生不同影響，結(jié)果不可預(yù)知。
為了解決這個(gè)問題，開發(fā)者可以人為的加入bot和top的依賴關(guān)系，讓指定運(yùn)算符先完成，如下圖所示：

人為的加入bot和top的依賴關(guān)系，讓指定運(yùn)算符先完成
如果需要保證讀取的值最新，需要新增下圖中虛線箭頭表示的依賴關(guān)系，即下圖中上方藍(lán)色圓圈依賴下方藍(lán)色圓圈的運(yùn)算完成，才能進(jìn)行計(jì)算。

加入依賴關(guān)系后，計(jì)算圖長這樣
條件分支
接下來看條件分支，即TensorFlow如何處理在這一節(jié)開頭提出來的那個(gè)例子？

TensorFlow提供了兩個(gè)條件控制OP，即tf.cond和tf.case
下面的代碼中，利用了tf.cond實(shí)現(xiàn)條件分支，在 a < b 為真，對(duì) out 求值會(huì)執(zhí)行 tf.add(3, 3)；否則，執(zhí)行 tf.square(3)。

使用tf.cond實(shí)現(xiàn)條件分支
上面這段代碼等價(jià)于：tf.cond(a < b, lambda: tf.add(3, 3), lambda: tf.sqaure(3))
然后生成的計(jì)算圖如下所示：

帶有條件控制流的計(jì)算圖
當(dāng)并列的分支比較多時(shí)，可以使用tf.case來處理，例如：

并列的條件分支>2個(gè)時(shí)，使用tf.case來控制
循環(huán)
TensorFlow提供了tf.while_loop來構(gòu)造循環(huán)塊，感覺和RNN類似的結(jié)構(gòu)有這個(gè)需求，例如：

tf.while_loop可以實(shí)現(xiàn)循環(huán)控制流解決RNN這種計(jì)算圖結(jié)構(gòu)的控制邏輯
下面的代碼實(shí)現(xiàn)了一個(gè)基礎(chǔ)的循環(huán)例子，即循環(huán)100次。

使用tf.while_loop在靜態(tài)圖中實(shí)現(xiàn)循環(huán)控制流
總的來說，TensorFlow應(yīng)該是首個(gè)將Control Flow引入到計(jì)算圖中的深度學(xué)習(xí)框架，不是像動(dòng)態(tài)圖框架那樣直接在Python層去做Control Flow，這方面必須給予一定的尊重。即使Pytorch目前在學(xué)術(shù)界已經(jīng)比TensorFlow更加流行，但基于TensorFlow演化的各種工業(yè)級(jí)項(xiàng)目仍然發(fā)揮著作用。
3. Pytorch中的Control Flow
在Pytorch這種動(dòng)態(tài)圖框架中，支持直接在Python端寫Control Flow，并且可以將這些控制邏輯放到計(jì)算圖中。這里以TorchScript為例，當(dāng)嘗試將Pytorch模型轉(zhuǎn)為TorchScript時(shí)，有兩種方式，一種是trace，另外一種是script。對(duì)于trace模式，適合Python層沒有Control Flow的計(jì)算圖，舉例如下：
#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() # 實(shí)例化模型
trace_module = torch.jit.trace(model,torch.rand(1,1,224,224))
print(trace_module.code) # 查看模型結(jié)構(gòu)
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模式則適用于計(jì)算圖在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.shapeif c ==1:x = self.conv1(x)else:x = self.conv2(x)return x

model = MyModule()

這樣寫會(huì)報(bào)錯(cuò)，因?yàn)橛锌刂屏?/h1>

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

此時(shí)應(yīng)該用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
然后來實(shí)驗(yàn)一下將上面帶有Control Flow的Module導(dǎo)出ONNX，這里以Pytorch官方文檔提供的一個(gè)帶循環(huán)的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)

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

可以看到直接導(dǎo)出Module，Python層的控制邏輯被丟掉（即for循環(huán)被完全展開），這是因?yàn)镻ytorch在導(dǎo)出ONNX的時(shí)候默認(rèn)使用了tracing機(jī)制
而當(dāng)使用script模式時(shí)，導(dǎo)出的ONNX就會(huì)保留Python層的Control Flow并將其轉(zhuǎn)換成ONNX中的Loop OP。示例代碼以及Netron可視化結(jié)果如下：
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. 總結(jié)
這篇文章介紹了一下深度學(xué)習(xí)中的Data Flow和Control Flow，然后介紹了一下將Pytorch模型轉(zhuǎn)為TorchScript的兩種模式，并探索了要將Pytorch的Python層的Control Flow轉(zhuǎn)換為ONNX應(yīng)該怎么做。
5. 參考文獻(xiàn)
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

總結(jié)

以上是生活随笔為你收集整理的深度学习编译器Data Flow和Control Flow的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。