摘要：是提供的一個(gè)非常具有前景的工具它能夠?qū)⒁徊糠终Z(yǔ)法的代碼轉(zhuǎn)譯成高效的圖表示代碼由于從開(kāi)始將會(huì)默認(rèn)使用動(dòng)態(tài)圖因此利用在理想情況下能讓我們實(shí)現(xiàn)用動(dòng)態(tài)圖寫方便靈活用靜態(tài)圖跑高效穩(wěn)定但是在使用的過(guò)程中如無(wú)意外肯定是會(huì)有意外的這篇文章就是指出一些和的奇

AutoGraph是TF提供的一個(gè)非常具有前景的工具, 它能夠?qū)⒁徊糠謕ython語(yǔ)法的代碼轉(zhuǎn)譯成高效的圖表示代碼. 由于從TF 2.0開(kāi)始,?TF將會(huì)默認(rèn)使用動(dòng)態(tài)圖(eager execution), 因此利用AutoGraph,?在理想情況下, 能讓我們實(shí)現(xiàn)用動(dòng)態(tài)圖寫(方便, 靈活), 用靜態(tài)圖跑(高效, 穩(wěn)定).

但是! 在使用的過(guò)程中, 如無(wú)意外肯定是會(huì)有意外的, 這篇文章就是指出一些AutoGraph和tf.function的奇怪的行為, 讓你更愉快地使用它們.

本文假設(shè)讀者具有一定的Python和TensorFlow的使用經(jīng)驗(yàn).

對(duì)tf1.X有經(jīng)驗(yàn)的讀者應(yīng)該不會(huì)對(duì)讓我們又愛(ài)又恨的計(jì)算圖(tf.Graph)和執(zhí)行會(huì)話(tf.Session)感到陌生, 一個(gè)常規(guī)的流程如下:

初始化一個(gè)計(jì)算圖并且將該計(jì)算圖設(shè)置為當(dāng)前scope下的默認(rèn)計(jì)算圖

用TF API設(shè)計(jì)計(jì)算圖(比如: y=tf.matmul(a, x) + b)

提前界定好參數(shù)共享并劃分相應(yīng)的參數(shù)scope

創(chuàng)建并配置好tf.Session

將計(jì)算圖傳給tf.Session

初始化參數(shù)

用tf.Session.run來(lái)執(zhí)行計(jì)算圖的節(jié)點(diǎn), 被執(zhí)行的節(jié)點(diǎn)會(huì)反向追蹤所有依賴的需要執(zhí)行的節(jié)點(diǎn)并執(zhí)行計(jì)算.

以下是上述過(guò)程的一個(gè)代碼例子:

g = tf.Graph() #初始化計(jì)算圖
with g.as_default(): # 設(shè)置為默認(rèn)計(jì)算圖
    a = tf.constant([[10,10],[11.,1.]]) 
    x = tf.constant([[1.,0.],[0.,1.]])
    b = tf.Variable(12.)
    y = tf.matmul(a, x) + b # 描述計(jì)算圖
    init_op = tf.global_variables_initializer() # 待執(zhí)行節(jié)點(diǎn)

with tf.Session() as sess: # 配置會(huì)話
    sess.run(init_op) # 執(zhí)行節(jié)點(diǎn)
    print(sess.run(y)) # 輸出結(jié)果

在TF 2.0中, 由于默認(rèn)為動(dòng)態(tài)圖, 計(jì)算會(huì)直接被執(zhí)行, 也就是說(shuō), 我們不需要

定義計(jì)算圖

會(huì)話執(zhí)行

參數(shù)初始化

用scope定義參數(shù)分享

用tf.control_dependencies來(lái)聲明節(jié)點(diǎn)的非直接依賴

我們可以像寫普通python代碼(or pytorch)一樣, 寫了就執(zhí)行:

a = tf.constant([[10,10],[11.,1.]])
x = tf.constant([[1.,0.],[0.,1.]])
b = tf.Variable(12.)
y = tf.matmul(a, x) + b
print(y.numpy())

一般來(lái)說(shuō), eager代碼會(huì)比執(zhí)行相同操作的靜態(tài)圖代碼的效率低, 因?yàn)楹芏嘤?jì)算圖優(yōu)化的方法只能用在數(shù)據(jù)流圖上.

如果想在TF 2.0上構(gòu)建傳統(tǒng)的計(jì)算圖, 我們就需要用到tf.function.

TF 2.0的其中一個(gè)重要改變就是去除tf.Session(此處應(yīng)有掌聲). 這個(gè)改變會(huì)迫使用戶用更好的方式來(lái)組織代碼: 不用再用讓人糾結(jié)的tf.Session來(lái)執(zhí)行代碼, 就是一個(gè)個(gè)python函數(shù), 加上一個(gè)簡(jiǎn)單的裝飾器.

在TF 2.0里面, 如果需要構(gòu)建計(jì)算圖, 我們只需要給python函數(shù)加上@tf.function的裝飾器.

上文提到靜態(tài)圖的執(zhí)行效率更高, 但是加速并不是一定的. 一般來(lái)說(shuō), 計(jì)算圖越復(fù)雜, 加速效果越明顯. 對(duì)于復(fù)雜的計(jì)算圖, 比如訓(xùn)練深度學(xué)習(xí)模型, 獲得的加速是巨大的. (譯者注: 個(gè)人感覺(jué)還是要結(jié)合實(shí)際來(lái)看, 如果某一部分的計(jì)算既有復(fù)雜的計(jì)算圖, 而計(jì)算圖的復(fù)雜性又帶來(lái)了額外的內(nèi)存消耗
或者計(jì)算量, 那么加速會(huì)比較明顯, 但是很多時(shí)候, 比如一般的CNN模型, 主要計(jì)算量并不在于圖的復(fù)雜性, 而在于卷積、矩陣乘法等操作, 加速并不會(huì)很明顯. 此處想法有待驗(yàn)證)

這個(gè)自動(dòng)將python代碼轉(zhuǎn)成圖表示代碼的工具就叫做AutoGraph.

在TF 2.0中, 如果一個(gè)函數(shù)被@tf.function裝飾了, 那么AutoGraph將會(huì)被自動(dòng)調(diào)用, 從而將python函數(shù)轉(zhuǎn)換成可執(zhí)行的圖表示.

在第一次調(diào)用被@tf.function裝飾的函數(shù)時(shí), 下列事情將會(huì)發(fā)生:

該函數(shù)被執(zhí)行并跟蹤。和Tensorflow 1.x類似, Eager會(huì)在這個(gè)函數(shù)中被禁用，因此每個(gè)tf.API只會(huì)定義一個(gè)生成tf.Tensor輸出的節(jié)點(diǎn)

AutoGraph用于檢測(cè)可以轉(zhuǎn)換為等效圖表示的Python操作（while→tf.while，for→tf.while，if→tf.cond，assert→tf.assert...)

為了保留執(zhí)行順序，在每個(gè)語(yǔ)句之后自動(dòng)添加tf.control_dependencies，以便在執(zhí)行第i+1行時(shí)確保第i行已經(jīng)被執(zhí)行. 至此計(jì)算圖已經(jīng)確定

根據(jù)函數(shù)名稱和輸入?yún)?shù)，創(chuàng)建唯一ID并將其與定義好的計(jì)算圖相關(guān)聯(lián)。計(jì)算圖被緩存到一個(gè)映射表中：map [id] = graph

如果ID配對(duì)上了，之后的函數(shù)調(diào)用都會(huì)直接使用該計(jì)算圖

下一節(jié)將會(huì)具體闡述如何將TF 1.X代碼塊分別改寫到eager和計(jì)算圖版本.

要使用tf.function, 第一步需要先將TF 1.X的設(shè)計(jì)計(jì)算圖的代碼放進(jìn)python函數(shù)里面.

def f():
    a = tf.constant([[10,10],[11.,1.]])
    x = tf.constant([[1.,0.],[0.,1.]])
    b = tf.Variable(12.)
    y = tf.matmul(a, x) + b
    return y

應(yīng)為TF 2.0默認(rèn)是eager的, 我們可以直接執(zhí)行該函數(shù)(不需要tf.Session):

print(f().numpy())

我們就會(huì)得到輸出:

[[22. 22.]
 [23. 13.]]

我們可以直接用@tf.function來(lái)裝飾函數(shù)f, 我們?cè)谠瓉?lái)f的基礎(chǔ)上加上宇宙第一的debug大法: print來(lái)更好地看看究竟發(fā)生了什么.

@tf.function
def f():
    a = tf.constant([[10,10],[11.,1.]])
    x = tf.constant([[1.,0.],[0.,1.]])
    b = tf.Variable(12.)
    y = tf.matmul(a, x) + b
    print("PRINT: ", y)
    tf.print("TF-PRINT: ", y)
    return y

f()

所以發(fā)生了什么呢?

@tf.function將函數(shù)f包進(jìn)了tensorflow.python.eager.def_function.Function這個(gè)對(duì)象, 函數(shù)f被賦予到了這個(gè)對(duì)象的.python_function屬性.

當(dāng)f()被執(zhí)行的時(shí)候, 計(jì)算圖會(huì)同時(shí)被構(gòu)建, 但是計(jì)算不會(huì)執(zhí)行, 因此我們會(huì)得到以下結(jié)果, tf.的操作不會(huì)被執(zhí)行:

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

最終, 你會(huì)看到代碼會(huì)執(zhí)行失敗:

ValueError: tf.function-decorated function tried to create variables on non-first call.

在?RFC: Functions, not Session里面有個(gè)非常明確的指示:

State (like tf.Variable objects) are only created the first time the function f is called. 狀態(tài)(比如tf.Variable) 只會(huì)在函數(shù)被第一次調(diào)用時(shí)創(chuàng)建.

但是?Alexandre Passos指出, 在函數(shù)轉(zhuǎn)換成圖表示時(shí), 我們沒(méi)有辦法確定tf.function調(diào)用了多少次函數(shù), 因此我們?cè)诘谝淮握{(diào)用函數(shù)f時(shí), 在圖構(gòu)建的過(guò)程中, 可能會(huì)被執(zhí)行了多次, 這就導(dǎo)致了上述錯(cuò)誤.

造成這個(gè)錯(cuò)誤的根源在于同樣的命令在動(dòng)態(tài)圖和靜態(tài)圖中的不一致性. 在動(dòng)態(tài)圖中, tf.Variable時(shí)一個(gè)普通的python變量, 超出了其作用域范圍就會(huì)被銷毀. 而在靜態(tài)圖中, tf.Variable則是計(jì)算圖中一個(gè)持續(xù)存在的節(jié)點(diǎn), 不受python的作用域的影響. 因此, 這是使用tf.function的第一個(gè)教訓(xùn):

將一個(gè)在動(dòng)態(tài)圖中可行的函數(shù)轉(zhuǎn)換成靜態(tài)圖需要用靜態(tài)圖的方式思考該函數(shù)是否可行

那么我們可以怎樣去規(guī)避這個(gè)錯(cuò)誤呢?

將tf.Variable作為函數(shù)的參數(shù)傳入

將父作用域繼承tf.Variable

將tf.Variable作為類屬性來(lái)調(diào)用

這里指方法2和方法3. 顯然的, 我們推薦使用方法3:

class F():
    def __init__(self):
        self._b = None

    @tf.function
    def __call__(self):
        a = tf.constant([[10, 10], [11., 1.]])
        x = tf.constant([[1., 0.], [0., 1.]])
        if self._b is None:
            self._b = tf.Variable(12.)
        y = tf.matmul(a, x) + self._b
        print("PRINT: ", y)
        tf.print("TF-PRINT: ", y)
        return y

f = F()
f()

我們之后會(huì)看到, 我們并不能隨意地用tf.function來(lái)轉(zhuǎn)化eager的代碼并達(dá)到加速的目的, 我們需要想象一下轉(zhuǎn)化是怎么完成的, 在轉(zhuǎn)python的代碼到圖操作的時(shí)候究竟發(fā)生了什么, 這些轉(zhuǎn)化包含了什么黑魔法. 這里的例子比較簡(jiǎn)單, 我們會(huì)在接下來(lái)的文章中更深入的探討.

@tf.function
def f(b):
    a = tf.constant([[10,10],[11.,1.]])
    x = tf.constant([[1.,0.],[0.,1.]])
    y = tf.matmul(a, x) + b
    print("PRINT: ", y)
    tf.print("TF-PRINT: ", y)
    return y

b = tf.Variable(12.)
f(b)

上述函數(shù)會(huì)得到我們想要的結(jié)果, 另外, 作為參數(shù)被傳入的變量能夠在函數(shù)中直接更新, 而更新后的值會(huì)在函數(shù)外也適用. 下面的代碼會(huì)打印出1,2,3

a = tf.Variable(0)

@tf.function
def g(x):
    x.assign_add(1)
    return x

print(g(a))
print(g(a))
print(g(a))

我們可以用@tf.function裝飾器來(lái)將python代碼轉(zhuǎn)成圖表示代碼

我們不能在被裝飾函數(shù)中初始化tf.Variable

可以用變量作用域繼承(對(duì)象屬性)或者參數(shù)傳入的方法使用在函數(shù)外初始化的變量

在之后的部分我們會(huì)更加深入地探討輸入?yún)?shù)類型對(duì)效率的影響, 以及python操作的轉(zhuǎn)換細(xì)節(jié).

聲明: 本文翻譯自Paolo Galeone的博客, 已取得作者的同意, 如需轉(zhuǎn)載本文請(qǐng)聯(lián)系本人

Disclaimer: This is a translation of the article?Analyzing tf.function to discover AutoGraph strengths and subtleties?by Paolo Galeone.