TensorFlow發佈重要更新AutoGraph

選自Medium，做者：Alex Wiltschko、Dan Moldovan、Wolff Dobson，機器之心編輯部。

近日，谷歌發佈了一項新的 TensorFlow 工具「AutoGraph」，能將 print() 函數和其它 Python 代碼轉化爲純 TensorFlow 計算圖代碼。這一工具極大地增強了 TensorFlow 在調用純 Python 語句時的性能，開發者能夠輕鬆在 TensorFlow 上實現更好的模型性能。

項目地址：github.com/tensorflow/…

通常而言，在寫 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：github.com/tensorflow/…

Colab：colab.research.google.com/github/tens…

如下咱們使用循環和分支來測試「科拉茲猜測」。注意，考慮到多樣性，咱們將不使用 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 上，請提出你的建議吧！

提建議：github.com/tensorflow/…

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 核心模塊。

原文連接：medium.com/tensorflow/…