為TensorFlow引入即時執行

谷歌開發者_發表於2017-11-20

我們宣佈為 TensorFlow 引入即時執行。即時執行是一個由執行定義的命令式介面,在這個介面中,運算將在從 Python 呼叫時立即執行。這樣可以讓 TensorFlow 的入門變得更加簡單,並讓研發變得更加直觀。

即時執行的好處包括:

  • 快速除錯即時執行時錯誤以及與 Python 工具整合

  • 支援使用易用型 Python 控制流的動態模型

  • 為自定義和高階漸變提供強大支援

  • 適用於幾乎所有可用的 TensorFlow 運算

  • 即時執行目前以實驗性功能形式提供,因此,我們希望獲得來自社群的反饋,以便指引我們的工作方向。


為了更好地理解這一功能,我們來看一些程式碼。它的技術性非常強;熟悉 TensorFlow 會有所幫助。



使用即時執行

啟用即時執行後,運算將立即執行並將它們的值返回到 Python,無需 Session.run()。例如,要將兩個矩陣相乘,我們可以這樣編寫程式碼:

import tensorflow as tf
import tensorflow.contrib.eager as tfe

tfe.enable_eager_execution()

x = [[2.]]
m = tf.matmul(x, x)


使用 print 或者 Python 除錯程式檢查中間結果非常直接。

print(m)
# The 1x1 matrix [[4.]]


動態模型可以通過 Python 流控制構建。下面是一個使用 TensorFlow 的算數運算的考拉茲猜想示例:

a = tf.constant(12)
counter = 0
while not tf.equal(a, 1):
  if tf.equal(a % 2, 0):
    a = a / 2
  else:
    a = 3 * a + 1
  print(a)


在這裡,tf.constant(12) 張量物件的使用會將所有數學運算提升為張量運算,因此,所有返回值都是張量。



漸變

大多數 TensorFlow 使用者都對自動微分感興趣。因為每次呼叫都可能發生不同的運算,我們將所有正向運算記錄到磁帶中,然後在計算漸變時進行倒放。在計算完漸變後,我們將捨棄磁帶。

如果您熟悉 autograd 軟體包,就會發現 API 將非常相似。例如:

def square(x):
  return tf.multiply(x, x)

grad = tfe.gradients_function(square)

print(square(3.))    # [9.]
print(grad(3.))      # [6.]

gradients_function 呼叫將 Python 函式 square() 用作一個引數並返回一個 Python 可呼叫物件,這個物件可以根據輸入計算 square() 的偏導數。因此,要獲得 square() 在輸入為 3.0 時的導數,請呼叫 grad(3.0),結果為 6。

相同的 gradients_function 呼叫可用於獲得平方的二階導數:

gradgrad = tfe.gradients_function(lambda x: grad(x)[0])

print(gradgrad(3.))  # [2.]

如前文所述,控制流會引起不同的運算,如下面的示例中所示。

def abs(x):
  return x if x > 0. else -x

grad = tfe.gradients_function(abs)

print(grad(2.0))  # [1.]
print(grad(-2.0)) # [-1.]



自定義漸變

使用者可能希望為某個運算或函式定義自定義漸變。這可能非常有用,原因之一是它為一系列運算提供了一種更高效、數值更穩定的漸變。

下面的示例演示了自定義漸變的使用。我們先來看一下 log(1 + ex) 函式,它通常用於計算交叉熵和對數似然性。

def log1pexp(x):
  return tf.log(1 + tf.exp(x))
grad_log1pexp = tfe.gradients_function(log1pexp)

# The gradient computation works fine at x = 0.
print(grad_log1pexp(0.))
# [0.5]
# However it returns a `nan` at x = 100 due to numerical instability.
print(grad_log1pexp(100.))
# [nan]

我們可以為上面的函式使用自定義漸變,從分析角度簡化漸變表示式。注意下面的漸變函式實現重用了在正向傳遞期間計算的表示式 (tf.exp(x)),通過避免冗餘計算提高了漸變計算的效率。

@tfe.custom_gradient
def log1pexp(x):
  e = tf.exp(x)
  def grad(dy):
    return dy * (1 - 1 / (1 + e))
  return tf.log(1 + e), grad
grad_log1pexp = tfe.gradients_function(log1pexp)

# Gradient at x = 0 works as before.
print(grad_log1pexp(0.))
# [0.5]
# And now gradient computation at x=100 works as well.
print(grad_log1pexp(100.))
# [1.0]



構建模型

模型可以分成幾類。下面的模型類建立了一個(簡單的)兩層網路來對標準的 MNIST 手寫數字進行分類。

class MNISTModel(tfe.Network):
  def __init__(self):
    super(MNISTModel, self).__init__()
    self.layer1 = self.track_layer(tf.layers.Dense(units=10))
    self.layer2 = self.track_layer(tf.layers.Dense(units=10))
  def call(self, input):
    """Actually runs the model."""
    result = self.layer1(input)
    result = self.layer2(result)
    return result

由於 tf.layer 建立幷包含模型引數(變數),我們建議在 tf.layer 中使用類,而不是函式。變數生命週期與層物件的生命週期關聯,因此,務必跟蹤它們。

為什麼要使用 tfe.Network?一個 Network 是一個層容器並且本身是一個 tf.layer.Layer,這樣,可以將 Network 物件嵌入到其他 Network 物件中。它還包含可以協助檢查、儲存和恢復的實用程式。

即使沒有訓練模型,我們也可以命令式地呼叫它並檢查輸出:

# Let's make up a blank input image
model = MNISTModel()
batch = tf.zeros([1, 1, 784])
print(batch.shape)
# (1, 1, 784)
result = model(batch)
print(result)
# tf.Tensor([[[ 0.  0., ...., 0.]]], shape=(1, 1, 10), dtype=float32)

請注意,我們不需要任何佔位符或會話。在我們第一次傳入輸入時,層引數的大小就會設定好。

為了訓練任何模型,我們需要定義一個損失函式來優化、計算漸變,並使用一個優化器來更新變數。首先,下面是一個損失函式:

def loss_function(model, x, y):
  y_ = model(x)
  return tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=y_)

然後,我們的訓練迴圈如下所示:

optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
for (x, y) in tfe.Iterator(dataset):
  grads = tfe.implicit_gradients(loss_function)(model, x, y)
  optimizer.apply_gradients(grads)


implicit_gradients() 將相對於計算期間使用的所有 TensorFlow 變數計算 loss_function 的導數。

我們可以按照往常處理 TensorFlow 的相同方式將計算轉移到 GPU 中:

with tf.device("/gpu:0"):
  for (x, y) in tfe.Iterator(dataset):
    optimizer.minimize(lambda: loss_function(model, x, y))


(注:我們簡化了損失函式的儲存並直接呼叫 optimizer.minimize,但是,您也可以使用上面的 apply_gradients() 函式;它們是等效的。)



結合使用即時執行與圖表

即時執行讓開發和除錯更具互動性,但是,TensorFlow 圖表在分散式訓練、效能優化和生產部署方面也有很多優勢。

在啟用即時執行時執行運算的相同程式碼將構建一個圖表來說明未啟用即時執行時的計算。要將您的模型轉換成圖表,只需在未啟用即時執行的新 Python 會話中執行相同的程式碼,如我們在 MNIST 示例中看到的一樣。模型變數的值可儲存並從檢查點恢復,這讓我們可以在即時(命令式)與圖表(宣告式)程式設計之間輕鬆切換。因此,可以輕鬆匯出在啟用即時執行時開發的模型,用於生產部署。

在不久的將來,我們會提供實用程式,讓您可以選擇性地將部分模型轉換成圖表。這樣,您可以融合部分計算(例如自定義 RNN 細胞的內部結構)來實現高效能,同時保持即時執行的靈活性和可讀性。



怎樣改寫我的程式碼?

即時執行的使用方法對當前 TensorFlow 使用者來說應當是直觀的。目前只有少量特定於即時執行的 API;大多數的現有 API 和運算都適用於即時執行。下面是一些需要注意的方面:

  • 一如我們通常對 TensorFlow 的建議,如果您還沒有從佇列轉到使用 tf.data 來處理輸入,我們建議您立即行動。TensorFlow 的使用更加簡單,並且速度通常更快。

  • 使用以物件為導向的層(例如 tf.layer.Conv2D())或 Keras 層;這些層可以顯式儲存變數。

  • 對於大多數模型,您可以編寫程式碼,讓它同時適用於即時執行和圖表構建。也有一些例外,例如,使用 Python 控制流根據輸入改變計算的動態模型就不行。

  • 一旦您呼叫 tfe.enable_eager_execution(),它將無法關閉。要獲取圖表行為,請啟動新的 Python 會話。



入門和未來

這仍是一個預覽版,因此,您可能會遇到一些不盡人意的地方。如果您想立即開始體驗,請執行以下操作:

  • 安裝 Nightly 版本的 TensorFlow。

  • 檢視 README(包括已知問題)

  • 在即時執行使用者指南中獲取詳細的說明

  • 瀏覽 GitHub 中的即時執行示例

  • 關注 changelog 獲取更新。


關於即時執行,我們還有很多要與大家分享的,我們非常高興,或者說非常渴望您立即試用這一功能。熱烈歡迎大家提供反饋。


0?wx_fmt=gif

相關文章