TensorFlow新功能「AutoGraph」：將Python轉換為計算圖

伊瓢 編譯自 TensorFlow博客

量子位 報道 | 公眾號 QbitAI

昨天，TensorFlow推出了一個新功能「AutoGraph」，可以將Python代碼（包括控制流print()和其他Python原生特性）轉換為TensorFlow的計算圖（Graph）代碼。

在不使用Eager Execution的情況下，編寫TensorFlow代碼需要進行一些元編程——先編寫一個創建計算圖的程序，然後稍後執行該程序。這就比較麻煩了，尤其是對新手來說。一些特別棘手的情況涉及更複雜的模型，例如使用if和while，或者是有副作用的模型print()，或接受結構化輸入。

AutoGraph

所以為什麼要轉化成計算圖呢？計算圖可以做各種優化，例如刪除常見的子表達式和內核融合。

而且計算圖讓分散式訓練和部署到各種環境更為容易，因為它們形成了獨立於平台的計算模型。這對於多個GPU或TPU上的分散式訓練尤為重要，或者通過TensorFlow Lite在移動或物聯網等其他平台上分發模型。

下面這個例子你可能想要添加到計算圖中：

如果照常使用Eager Execution，它完全可以「正常工作」，但是由於Python解釋器開銷或者沒有進行程序優化，它可能執行的很慢。

所以，計算圖執行需要一個前提條件：用類似tf.cond()的結構重寫它，就是可能會比較無聊，並且難以實現。

現在，AutoGraph可以自動完成這個轉換的過程，這樣可以既簡單又能獲取基於計算圖執行的性能優勢。

在這個例子中，我們可以用autograph.convert()來裝飾函數，AutoGraph將自動生成計算圖就緒代碼。

使用AutoGraph，原來的這段代碼：

在裝飾器的作用下變成下面這段：

現在，就可以直接調用代碼，就像調用TensorFlow op一樣：

就這樣，AutoGraph填補了eager execution和計算圖之間的空白，AutoGraph可以把eager-style的Python代碼轉換為graph-generating的代碼。

AutoGraph不僅僅是一組有用的宏; 它使用源代碼轉換來Python的任何部分，包括控制流、函數應用程序和賦值、生成樣板代碼、以及重構常用的Python代碼使其容易轉換為計算圖。

另外，不管使用什麼編譯器，都需要保證報錯信息可讀。為此，AutoGraph設置了創建錯誤消息和堆棧跟蹤，可以幫你找到代碼中的錯誤源，而不是僅僅是引用錯誤代碼。

可運行的例子

這裡TensorFlow官方展示了一個用循環和分支檢查Collatz猜想的例子，用AutoGraph的 .to_graph()函數將其轉換為計算圖：

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

AutoGraph允許您將元素追加到循環內的數組中，可以通過使用一些AutoGraph助手，比如set_element_type和stack來實現。

我們還支持像break、continue、print、assert等這些結構，轉換後，該部分Python代碼中的assert將轉換為tf.Assert來表示計算圖。

能夠輕鬆地添加循環，控制流程以及更多計算圖，意味著可以輕鬆地將訓練循環移動到計算圖中。另外的一個例子是採用RNN訓練循環並通過一次sess.run()調用執行它。在需要將整個訓練循環傳遞給加速器而不是通過CPU控制器管理訓練的情況下，這可能很有用。

如果想看更多的例子，本文文末有TensorFlow官方github例子傳送門。

轉換為計算圖 vs Eager Execution

雖然Eager Execution很有用，但是計算圖更快。雖然基準很複雜（並且依賴於程序和硬體本身），但在這個例子（鏈接：http://t.cn/RgCsKOe）里，我們看到了從eager execution切換到AutoGraph代碼後速度大大提升。

最後，AutoGraph可以讓你在GPU或者雲端TPU等加速器硬體上使用動態模型或者重度控制流模型，用大數據訓練大型模型時這是必須的。

AutoGraph和Eager Execution

將來，AutoGraph將和defun無縫集成，以在eager-style的代碼中生成計算圖。屆時，你可以通過把eager代碼轉換為計算圖片段來使用AutoGraph加速。

然鵝還是試驗工具

雖然AutoGraph看起來很好用，不過TensorFlow官方博客的最後還是說，它還是contrib里的實驗工具，不過，官方會儘快將其轉移到核心TensorFlow中。

傳送門

AutoGraph：https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/autograph

代碼示例：https://github.com/tensorflow/models/blob/master/samples/core/guide/autograph.ipynb

—完—

活動報名

誠摯招聘

量子位正在招募編輯/記者，工作地點在北京中關村。期待有才氣、有熱情的同學加入我們！相關細節，請在量子位公眾號(QbitAI)對話界面，回復「招聘」兩個字。

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