選自Medium，作者：Alex Wiltschko、Dan Moldovan、Wolff Dobson，機器之心編輯部。
近日，谷歌發佈了一項新的 TensorFlow 工具「AutoGraph」，能將 print() 函數和其它 Python 代碼轉化為純 TensorFlow 計算圖代碼。這一工具極大地加強了 TensorFlow 在調用純 Python 語句時的性能，開發者可以輕鬆在 TensorFlow 上實現更好的模型性能。
項目地址：https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/autograph
一般而言，在寫 TensorFlow 代碼時，我們需要構建整個算法的計算圖，或者規劃所有數據流的計算過程，然後再投入數據並快速執行整個或局部計算圖。當然因為當前 PyTorch 和 Keras 等動態計算圖的流行，TensorFlow 也發佈了 Eager Execution，它可以幫助用戶自動構建計算圖。但一般的 TensorFlow 還是常使用靜態計算圖的方式，因為它的構建邏輯與部署都非常有優勢。
然而對於入門開發者而言，理解靜態計算圖是比較困難的，因此很容易引起開發者的困惑。尤其是在一些涉及更復雜模型場景中，例如使用 if 和 while 等 Python 語句，或使用 print() 與接受結構化輸入等，它們都會引起我們對計算圖的困惑。
所以為什麼 TensorFlow 需要使用計算圖呢？計算圖允許各種各樣的優化，例如移除公共的子表達式和內核融合等。此外，計算圖簡化了分佈式訓練和部署時的環境配置，因此它們可被視為一種獨立於平臺的模型計算形式。這一特性對於在多 GPU 或 TPU 上的分佈式訓練極其重要，當然基於 TensorFlow Lite 在移動端和 IoT 上部署模型也非常重要。
以下是一個非常簡單的操作示例：
def huber_loss(a):
if tf.abs(a) <= delta:
loss = a * a / 2
else:
loss = delta * (tf.abs(a) - delta / 2)
return loss
使用 Eager Execution，這只是「正確運行」而已，但是此類操作可能會比較慢，因為 Python 解釋器眾所周知在實現地比較慢，且需要的計算比較複雜，這會令它錯過許多程序優化的機會。
為了給圖執行做好準備，你需要重寫代碼，使用 tf.cond() 等語句，但是這很繁瑣且難以實現。AutoGraph 可以自動完成該轉換，保持 Eager 編程的簡易性，同時還提升了計算圖執行的性能。
在該示例中，我們可以使用 autograph.convert() 佈置我們的函數，AutoGraph 將自動生成圖可用的代碼。
使用 AutoGraph，由於 decorator，下列代碼：
@autograph.convert()
def huber_loss(a):
if tf.abs(a) <= delta:
loss = a * a / 2
else:
loss = delta * (tf.abs(a) - delta / 2)
return loss
在執行時變成如下代碼：
def tf__huber_loss(a):
with tf.name_scope('huber_loss'):
def if_true():
with tf.name_scope('if_true'):
loss = a * a / 2
return loss,
def if_false():
with tf.name_scope('if_false'):
loss = delta * (tf.abs(a) - delta / 2)
return loss,
loss = ag__.utils.run_cond(tf.less_equal(tf.abs(a), delta), if_true,
if_false)
return loss
接下來，你可以調用你的代碼，就像使用一般的 TensorFlow op 一樣：
with tf.Graph().as_default():
x_tensor = tf.constant(9.0)
# The converted function works like a regular op: tensors in, tensors out.
huber_loss_tensor = huber_loss(x_tensor)
with tf.Session() as sess:
print('TensorFlow result: %2.2f\n' % sess.run(huber_loss_tensor))
如你所見，AutoGraph 連接起 Eager execution 和 Graph。AutoGraph 使用 Eager-style 的 Python 代碼，然後將其轉換成圖生成代碼。
AutoGraph 不只是有用宏命令的集合，它還可以使用源代碼轉換來覆寫 Python 語言的任意部分，包括控制流、函數應用和分配，生成樣板代碼，重構慣用 Python，以使轉換成圖的過程變得簡單。
使用任意編譯器，都會對錯誤信息可讀性產生擔憂；為此，AutoGraph 可以創建錯誤信息，並堆疊揭示原始源代碼中錯誤來源的多個軌跡，而不是僅僅顯示生成代碼的 reference。
可運行示例
那麼，AutoGraph 可以為我們做什麼呢？以下有一些示例代碼，它們可以直接轉換為圖代碼而不需要任何的改寫。如果你想實際運行這些操作，谷歌在這個 GitHub 的 Colab 中提供了一個 notebook 可供使用。
GitHub：https://github.com/tensorflow/models/blob/master/samples/core/guide/autograph.ipynb
Colab：https://colab.research.google.com/github/tensorflow/models/blob/master/samples/core/guide/autograph.ipynb
以下我們使用循環和分支來測試「科拉茲猜想」。注意，考慮到多樣性，我們將不使用 decorator，而使用 AutoGraph 的.to_graph() 函數將其轉換為圖。
def collatz(a):
counter = 0
while a != 1:
if a % 2 == 0:
a = a // 2
else:
a = 3 * a + 1
counter = counter + 1
return counter
graph_mode_collatz = autograph.to_graph(collatz)
# The code is human-readable, too
print(autograph.to_code(collatz))
collatz_tensor = graph_mode_collatz(tf.constant(n))
AutoGraph 可以支持任意的嵌套控制流，例如：
def f(n):
if n >= 0:
while n < 5:
n += 1
print(n)
return n
AutoGraph 允許你在循環中添加元素到數組中。為了讓其工作，我們使用一些 AutoGraph 輔助工具，set_element_type 和 stack。
def f(n):
z = []
# We ask you to tell us the element dtype of the list
autograph.set_element_type(z, tf.int32)
for i in range(n):
z.append(i)
# when you're done with the list, stack it
# (this is just like np.stack)
return autograph.stack(z)
我們還支持 break、continue，甚至 print 和 assert 等語句。當轉換完成後，這個片段的 Python assert 使用合適的 tf.Assert 將其轉換為 TensorFlow 計算圖。
def f(x):
assert x != 0, 'Do not pass zero!'
return x * x
具備輕易地添加循環、控制流等到圖上的能力意味著可以很容易將訓練循環轉移到圖中。可以在這個 Colab 的 notebook 中找到一個示例，其中使用了一個 RNN 訓練循環，並用一個 sess.run() 調用來執行它。當你需要傳遞一個完整的訓練循環到加速器時，這很有用，比通過 CPU 控制器管理訓練過程更好。
AutoGraph 打開了構建和訓練模型的新思路。谷歌在未來將基於開發者社區建議嘗試添加更多的功能到 AutoGraph 上，請提出你的建議吧！
提建議：https://github.com/tensorflow/tensorflow/issues
Graph Performance 對比 Eager Execution
Eager Execution 相當合用，但圖更快。儘管對比基準較為複雜（由應用以及硬件配置決定），但在一些簡單示例中我們可以看到，當從 Eager 轉換到 AutoGraph 代碼時有極大的加速，使用了大量 if 和 while 等語句。
最終，AutoGraph 讓你可以在 GPU 和 Cloud TPU 這樣的加速器硬件上使用動態和流控制極嚴模型，這對在大量數據上訓練大型模型非常有幫助。
AutoGraph 和 Eager Execution
雖然使用 Eager Execution，你也能通過 tf.contrib.eager.defun 對部分代碼根據計算圖執行。但這需要你使用 tf.cond() 這樣計算圖類的 TensorFlow ops。未來，AutoGraph 將無縫與 defun 融合，讓你用簡單的 eager-style Python 編寫圖代碼。當成為現實時，通過選擇性的把 eager 代碼轉換到圖分段，你就可以期待使用 AutoGraph 加速熱點了。
結論
AutoGraph 能夠讓你輕鬆的建立在 TensorFlow 圖中輕鬆運行的直觀性、複雜模型。這是目前在 contrib 中運行的實驗性工具，但我們期望能夠儘快把它加入到 TensorFlow 核心模塊。
閱讀更多 機器之心 的文章
