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

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

那么，我們為什么需要圖呢？圖允許各種優化，如移除常見的子表達式和融合內核。此外，由于圖形成了一種獨立于平臺的計算模型，可簡化分布式訓練和針對各種環境的部署。這對于多個 GPU 或 TPU 上的分布式訓練或者通過TensorFlow Lite在移動或物聯網等其他平臺上分發模型尤為重要。

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

1def huber_loss(a):

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

3loss = a * a / 2

4else:

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

6return loss

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

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

在本例中，我們可以使用 autograph.convert() 裝飾函數，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

在執行時變為以下代碼。

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 運算進行調用：

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 填補了即時執行與圖之間的差距。AutoGraph 接收即時風格的 Python 代碼并將其轉換為生成圖的代碼。

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

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

可運行示例

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

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

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

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 允許在循環內向數組追加元素。為此，我們將使用一些 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 等語句。轉換后，此代碼段的 Python assert 轉換為使用適當的 tf.Assert 的圖。

1def f(x):

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

3return x * x

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

注：筆記本鏈接

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

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

圖性能與即時執行

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

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

我們剛剛開始著手分析性能。如果您發現某個圖語句的運行速度低于預期，請提交問題！

AutoGraph 和即時執行

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

結論

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

跟我們分享您的 AutoGraph 使用體驗！如果您有任何反饋、建議或想法，請提交問題并向TensorFlow 開發者群組發送消息。