今天，我們將介紹一項(xiàng)名為 “AutoGraph” 的 TensorFlow 新功能。AutoGraph 可以將 Python 代碼（包括控制流、print() 和其他 Python 原生功能）轉(zhuǎn)換為純 TensorFlow 圖代碼。

要在不使用即時(shí)執(zhí)行的情況下編寫(xiě) TensorFlow 代碼，您需要進(jìn)行一些元編程，即編寫(xiě)一個(gè)用于創(chuàng)建圖的程序，然后再執(zhí)行該圖。這種方式可能會(huì)令人困惑，特別是對(duì)于新開(kāi)發(fā)者而言。一些特別棘手的情況涉及更復(fù)雜的模型，例如使用 if 和 while 的模型，或者具有 print() 等副作用或接受結(jié)構(gòu)化輸入的模型。

那么，我們?yōu)槭裁葱枰獔D呢？圖允許各種優(yōu)化，如移除常見(jiàn)的子表達(dá)式和融合內(nèi)核。此外，由于圖形成了一種獨(dú)立于平臺(tái)的計(jì)算模型，可簡(jiǎn)化分布式訓(xùn)練和針對(duì)各種環(huán)境的部署。這對(duì)于多個(gè) GPU 或 TPU 上的分布式訓(xùn)練或者通過(guò)TensorFlow Lite在移動(dòng)或物聯(lián)網(wǎng)等其他平臺(tái)上分發(fā)模型尤為重要。

下面這個(gè)簡(jiǎn)單示例顯示了您可能要添加到圖的運(yùn)算：

1def huber_loss(a):

2if tf.abs(a) <= delta: ? ?

3loss = a * a / 2

4else:

5loss = delta * (tf.abs(a) - delta / 2)

6return loss

如果使用即時(shí)執(zhí)行，也可以滿(mǎn)足要求，但是由于 Python 解釋器開(kāi)銷(xiāo)或錯(cuò)過(guò)程序優(yōu)化機(jī)會(huì)，這樣可能會(huì)很慢。

要準(zhǔn)備好執(zhí)行圖，您需要進(jìn)行重寫(xiě)以使用類(lèi)似 tf.cond() 的語(yǔ)句，但這可能比較繁瑣并且難以實(shí)現(xiàn)。AutoGraph 可以為您自動(dòng)執(zhí)行此轉(zhuǎn)換，既保持了即時(shí)編程的簡(jiǎn)易性，同時(shí)又獲得了圖執(zhí)行的性能優(yōu)勢(shì)。

在本例中，我們可以使用 autograph.convert() 裝飾函數(shù)，AutoGraph 將自動(dòng)生成可生成圖的代碼。

使用 AutoGraph 時(shí)，由于裝飾器的原因，以下代碼：

1@autograph.convert()

2def huber_loss(a):

3if tf.abs(a) <= delta: ? ?

4loss = a * a / 2

5else:

6loss = delta * (tf.abs(a) - delta / 2)

7return loss

在執(zhí)行時(shí)變?yōu)橐韵麓a。

1def tf__huber_loss(a):

2with tf.name_scope('huber_loss'):

3

4def if_true():

5with tf.name_scope('if_true'):

6loss = a * a / 2

7 return loss,

8

9def if_false():

10with tf.name_scope('if_false'):

11 loss = delta * (tf.abs(a) - delta / 2)

12 return loss,

13 loss = ag__.utils.run_cond(tf.less_equal(tf.abs(a), delta), if_true,

14if_false)

15 return loss

然后，您可以將這些代碼視為 TensorFlow 運(yùn)算進(jìn)行調(diào)用：

1with tf.Graph().as_default():

2x_tensor = tf.constant(9.0)

3

4# The converted function works like a regular op: tensors in, tensors out.

5huber_loss_tensor = huber_loss(x_tensor)

6

7with tf.Session() as sess:

8print('TensorFlow result: %2.2f\n' % sess.run(huber_loss_tensor))

如您所見(jiàn)，AutoGraph 填補(bǔ)了即時(shí)執(zhí)行與圖之間的差距。AutoGraph 接收即時(shí)風(fēng)格的 Python 代碼并將其轉(zhuǎn)換為生成圖的代碼。

AutoGraph 不僅僅是一組有用的宏，它還使用源代碼轉(zhuǎn)換來(lái)允許替換 Python 語(yǔ)言的任意部分，包括控制流、函數(shù)應(yīng)用和賦值、生成模板代碼，以及重構(gòu)慣用 Python 以便輕松轉(zhuǎn)換成圖。

使用任何編譯器，都會(huì)對(duì)錯(cuò)誤消息的可讀性產(chǎn)生擔(dān)憂(yōu)；為此，AutoGraph 設(shè)置為創(chuàng)建錯(cuò)誤消息和堆疊追蹤，以顯示原始源代碼中的錯(cuò)誤源，而不是僅顯示對(duì)生成代碼的引用。

可運(yùn)行示例

那么，AutoGraph 可以為我們做什么呢？以下是一些代碼示例，現(xiàn)在可以直接轉(zhuǎn)換為圖代碼而無(wú)需任何更改。如果您想實(shí)際運(yùn)行所有這些操作，我們提供了一個(gè)筆記本，您可以在Colab中打開(kāi)或在GitHub中查看。

注：GitHub 鏈接https://github.com/tensorflow/models/blob/master/samples/core/guide/autograph.ipynb

接下來(lái)，我們使用循環(huán)和分支來(lái)檢查Collatz 猜想。請(qǐng)注意，出于多樣性考慮，我們將不使用裝飾器，而使用 AutoGraph 的 .to_graph() 函數(shù)將其轉(zhuǎn)換為圖。

1def collatz(a):

2counter = 0

3while a != 1:

4if a % 2 == 0:

5a = a // 2

6else:

7a = 3 * a + 1

8counter = counter + 1

9return counter

10

11graph_mode_collatz = autograph.to_graph(collatz)

12# The code is human-readable, too

13print(autograph.to_code(collatz))

14

15collatz_tensor = graph_mode_collatz(tf.constant(n))

AutoGraph 可以支持任意的嵌套控制流，例如：

1def f(n):

2if n >= 0:

3while n < 5: ? ?

4n += 1

5print(n)

6return n

AutoGraph 允許在循環(huán)內(nèi)向數(shù)組追加元素。為此，我們將使用一些 AutoGraph 輔助工具，set_element_type 和 stack。

1def f(n):

2z = []

3# We ask you to tell us the element dtype of the list

4autograph.set_element_type(z, tf.int32)

5for i in range(n):

6z.append(i)

7# when you're done with the list, stack it

8# (this is just like np.stack)

9return autograph.stack(z)

我們還支持 break、continue，甚至 print 和 assert 等語(yǔ)句。轉(zhuǎn)換后，此代碼段的 Python assert 轉(zhuǎn)換為使用適當(dāng)?shù)?tf.Assert 的圖。

1def f(x):

2assert x != 0, 'Do not pass zero!'

3return x * x

能夠輕松向圖添加循環(huán)和控制流等，這意味著可以輕松將訓(xùn)練循環(huán)轉(zhuǎn)移到圖中。您可以在此筆記本中找到相關(guān)示例，其中，我們接受一個(gè) RNN 訓(xùn)練循環(huán)并用一個(gè) sess.run() 調(diào)用執(zhí)行。在需要將整個(gè)訓(xùn)練循環(huán)傳遞給加速器而不是通過(guò) CPU 控制器管理訓(xùn)練的情況下，這樣做十分有用。

注：筆記本鏈接

https://colab.research.google.com/github/tensorflow/models/blob/master/samples/core/guide/autograph.ipynb#scrollTo=4LfnJjm0Bm0B

AutoGraph 開(kāi)啟了構(gòu)建和訓(xùn)練模型的新思路。我們期待根據(jù)開(kāi)發(fā)者社區(qū)的建議為 AutoGraph 添加更多功能，因此歡迎您提交問(wèn)題并給出建議！

圖性能與即時(shí)執(zhí)行

即時(shí)執(zhí)行非常簡(jiǎn)單易用，但圖的速度通常要更快。雖然二者的比較基準(zhǔn)較為復(fù)雜（并且取決于應(yīng)用和硬件配置），但在這個(gè)簡(jiǎn)單的示例中，我們可以看到，從即時(shí)切換到大量使用 if 和 while 的 AutoGraph 代碼時(shí)，速度有了顯著提升。

最終，AutoGraph 允許您在加速器硬件（如 GPU 和Cloud TPU）上使用動(dòng)態(tài)和流控制較多的模型，這在基于大量數(shù)據(jù)訓(xùn)練大型模型時(shí)十分必要。

我們剛剛開(kāi)始著手分析性能。如果您發(fā)現(xiàn)某個(gè)圖語(yǔ)句的運(yùn)行速度低于預(yù)期，請(qǐng)?zhí)峤粏?wèn)題！

AutoGraph 和即時(shí)執(zhí)行

雖然使用即時(shí)執(zhí)行時(shí)，您仍然可以通過(guò)tf.contrib.eager.defun對(duì)部分代碼使用圖執(zhí)行，但這要求您使用 tf.cond() 等 TensorFlow 圖運(yùn)算。將來(lái)，AutoGraph 將與 defun 無(wú)縫集成，允許用簡(jiǎn)單的即時(shí)式 Python 語(yǔ)言編寫(xiě)圖代碼。實(shí)現(xiàn)此功能后，您可以通過(guò)選擇性地將即時(shí)代碼轉(zhuǎn)換為圖片段來(lái)使用 AutoGraph 加速熱點(diǎn)。

結(jié)論

AutoGraph 是一款新工具，可幫您輕松構(gòu)建能夠在 TensorFlow 圖中輕松運(yùn)行的直觀、復(fù)雜的模型。此工具目前為 contrib 中的實(shí)驗(yàn)性工具，但我們希望盡快將它加入到核心 TensorFlow 中。

跟我們分享您的 AutoGraph 使用體驗(yàn)！如果您有任何反饋、建議或想法，請(qǐng)?zhí)峤粏?wèn)題并向TensorFlow 開(kāi)發(fā)者群組發(fā)送消息。