一項名為“AutoGraph”的TensorFlow新功能

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

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

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

下面這個簡單示例顯示了您可能要添加到圖的運算：

1def huber_loss(a):

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

3loss = a * a / 2

4else:

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

6return loss

如果使用即時執(zhí)行，也可以滿足要求，但是由于 Python 解釋器開銷或錯過程序優(yōu)化機會，這樣可能會很慢。

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

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

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

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í)行時變?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 運算進行調(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))

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

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

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

可運行示例

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

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

接下來，我們使用循環(huán)和分支來檢查Collatz 猜想。請注意，出于多樣性考慮，我們將不使用裝飾器，而使用 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 等語句。轉(zhuǎn)換后，此代碼段的 Python assert 轉(zhuǎn)換為使用適當?shù)?tf.Assert 的圖。

1def f(x):

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

3return x * x

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

注：筆記本鏈接

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

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

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

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

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

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

AutoGraph 和即時執(zhí)行

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

結(jié)論

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

跟我們分享您的 AutoGraph 使用體驗！如果您有任何反饋、建議或想法，請?zhí)峤粏栴}并向TensorFlow 開發(fā)者群組發(fā)送消息。