TensorFlow新工具 | AutoGraph

前言

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

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

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

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 可供使用。

以下我們使用迴圈和分支來測試「科拉茲猜想」。注意，考慮到多樣性，我們將不使用 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 控制器管理訓練過程更好。

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中。

【本文轉載自： 深度學習與神經網路，作者：翻譯|Amusi，原文連結：https://mp.weixin.qq.com/s/XuTtsnwlO_0nZjK3BqUk0A】

來自 “ ITPUB部落格 ” ，連結：http://blog.itpub.net/31542119/viewspace-2200306/，如需轉載，請註明出處，否則將追究法律責任。