TensorFlow新工具 AutoGraph

重磅乾貨，第一時間送達

作者: Michael Tyka

翻譯: Amusi

來源: medium.com/tensorflow

前言

我們（TensorFlow官方）想告訴你一個名為「AutoGraph」的新TensorFlow功能。 AutoGraph將Python代碼（包括控制流，print（）和其他Python原生特徵）轉換為純TensorFlow圖形代碼。

為什麼我們需要圖（graph）？ 圖允許各種優化，例如刪除常見的子表達式和融合內核。此外，圖形簡化了分散式培訓和部署到各種環境，因為它們形成了獨立於平台的計算模型。這對於多個GPU或TPU上的分散式訓練尤為重要，或者通過TensorFlow Lite在移動或物聯網等其他平台上分發您的模型。

以下是一個非常簡單的操作示例：

使用 Eager Execution，這只是「正確運行」而已，但是此類操作可能會比較慢，因為 Python 解釋器眾所周知在實現地比較慢，且需要的計算比較複雜，這會令它錯過許多程序優化的機會。

為了給圖執行做好準備，你需要重寫代碼，使用 tf.cond() 等語句，但是這很繁瑣且難以實現。AutoGraph 可以自動完成該轉換，保持 Eager 編程的簡易性，同時還提升了計算圖執行的性能。

在該示例中，我們可以使用 autograph.convert() 布置我們的函數，AutoGraph 將自動生成圖可用的代碼。

使用 AutoGraph，由於 decorator，下列代碼：

在執行時變成如下代碼：

接下來，你可以調用你的代碼，就像使用一般的 TensorFlow op 一樣：

如你所見，AutoGraph 連接起 Eager execution 和 Graph。AutoGraph 使用 Eager-style 的 Python 代碼，然後將其轉換成圖生成代碼。

AutoGraph 不只是有用宏命令的集合，它還可以使用源代碼轉換來覆寫 Python 語言的任意部分，包括控制流、函數應用和分配，生成樣板代碼，重構慣用 Python，以使轉換成圖的過程變得簡單。

使用任意編譯器，都會對錯誤信息可讀性產生擔憂；為此，AutoGraph 可以創建錯誤信息，並堆疊揭示原始源代碼中錯誤來源的多個軌跡，而不是僅僅顯示生成代碼的 reference。

可運行示例

那麼，AutoGraph 可以為我們做什麼呢？以下有一些示例代碼，它們可以直接轉換為圖代碼而不需要任何的改寫。如果你想實際運行這些操作，谷歌在這個 GitHub 的 Colab 中提供了一個 notebook 可供使用。

以下我們使用循環和分支來測試「科拉茲猜想」。注意，考慮到多樣性，我們將不使用 decorator，而使用 AutoGraph 的.to_graph()函數將其轉換為圖。

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

AutoGraph 允許你在循環中添加元素到數組中。為了讓其工作，我們使用一些 AutoGraph 輔助工具，set_element_type 和 stack。

我們還支持 break、continue，甚至 print 和 assert 等語句。當轉換完成後，這個片段的 Python assert 使用合適的 tf.Assert 將其轉換為 TensorFlow 計算圖。

具備輕易地添加循環、控制流等到圖上的能力意味著可以很容易將訓練循環轉移到圖中。可以在這個 Colab 的 notebook 中找到一個示例，其中使用了一個 RNN 訓練循環，並用一個 sess.run() 調用來執行它。當你需要傳遞一個完整的訓練循環到加速器時，這很有用，比通過 CPU 控制器管理訓練過程更好。

Graph Performance 對比 Eager Execution

Eager Execution 相當合用，但圖更快。儘管對比基準較為複雜（由應用以及硬體配置決定），但在一些簡單示例中我們可以看到，當從 Eager 轉換到 AutoGraph 代碼時有極大的加速，使用了大量 if 和 while 等語句。

最終，AutoGraph 讓你可以在 GPU 和 Cloud TPU 這樣的加速器硬體上使用動態和流控制極嚴模型，這對在大量數據上訓練大型模型非常有幫助。

AutoGraph 和 Eager Execution

結論

AutoGraph是一款工具，可讓您輕鬆構建直觀，複雜的模型，在TensorFlow圖中輕鬆運行。 這是一個現在在contrib中的實驗工具，但我們希望儘快將其轉移到核心TensorFlow中。

告訴我們您使用AutoGraph的經歷！ 如果您有反饋，建議或想法，請提交問題並向TensorFlow開發人員小組發送消息。

致謝

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！