終於！TensorFlow引入了動態圖機制Eager Execution

機器之心發表於2017-11-01

PyTorch 的動態圖一直是 TensorFlow 使用者求之不得的功能，谷歌也一直試圖在 TensorFlow 中實現類似的功能。最近，Google Brain 團隊釋出了 Eager Execution，一個由執行定義的新介面，讓 TensorFlow 開發變得簡單許多。在工具推出後，谷歌開發人員 Yaroslav Bulatov 對它的效能與 PyTorch 做了橫向對比。

今天，我們為 TensorFlow 引入了「Eager Execution」，它是一個命令式、由執行定義的介面，一旦從 Python 被呼叫，其操作立即被執行。這使得入門 TensorFlow 變的更簡單，也使研發更直觀。

Eager Execution 的優點如下：

  • 快速除錯即刻的執行錯誤並通過 Python 工具進行整合
  • 藉助易於使用的 Python 控制流支援動態模型
  • 為自定義和高階梯度提供強大支援
  • 適用於幾乎所有可用的 TensorFlow 運算

Eager Execution 現在處於試用階段，因此我們希望得到來自社群的反饋，指導我們的方向。

為了更好地理解 Eager Execution，下面讓我們看一些程式碼。它很技術，熟悉 TensorFlow 會有所幫助。

使用 Eager Execution

當你啟動 Eager Execution 時，運算會即刻執行，無需 Session.run() 就可以把它們的值返回到 Python。比如，要想使兩個矩陣相乘，我們這樣寫程式碼：

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 算術操作的考拉茲猜想（Collatz conjecture）的一個示例：

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 使用者對自動微分（automatic differentiation）很感興趣。因為每次呼叫都有可能出現不同的運算，可以理解為我們把所有的正向運算錄到「磁帶」上，然後在計算梯度時進行「倒放」。梯度計算完成後，「磁帶」就沒用了。

如果你熟悉 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 callable，用於計算輸入的 square() 偏導數。因此，為了得到輸入為 3.0 時的 square() 導數，啟用 grad(3.0)，也就是 6。

同樣的 gradient_function 呼叫可用於計算 square() 的二階導數。

gradgrad = tfe.gradients_function(lambda x: grad(x)[0])
print(gradgrad(3.))  # [2.]

如前所述，控制流（control flow）會引起不同的運算，下面是一個示例：

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 + e^x)，它通常用於計算交叉熵和 log 似然。

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.layers 中的類別（而非函式），這是因為它們建立幷包含了模型引數（變數，variables）。變數的有效期和層物件的有效期緊密相關，因此需要對它們進行追蹤。

為什麼要使用 tfe.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)

注意我們在這裡不需要任何的佔位符或會話（session）。一旦資料被輸入，層的引數就被設定好了。

訓練任何模型都需要定義一個損失函式，計算梯度，並使用一個優化器更新引數。首先定義一個損失函式：


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 變數的導數。

我們可以按往常使用 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() 方法，它們是等價的。）

使用 Eager 和 Graphs

Eager execution 使開發和除錯互動性更強，但是 TensorFlow graph 在分散式訓練、效能優化和生產部署中也有很多優勢。

啟用 eager execution 時，執行運算的程式碼還可以構建一個描述 eager execution 未啟用時的計算圖。為了將模型轉換成圖，只需要在 eager execution 未啟用的 Python session 中執行同樣的程式碼。示例：https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples/mnist。我們可以從檢查點儲存和修復模型變數值，這允許我們在 eager（命令式）和 graph（宣告式）程式設計之間輕鬆轉換。這樣，啟用 eager execution 開發出的模型可以輕鬆匯出到生產部署中。

在不久的將來，我們將提供工具，可以選擇性地將模型的某些部分轉換成 graph。用這種方式，你就可以融合部分計算（如自定義 RNN 細胞的內部）實現高效能，同時還能保持 eager execution 的靈活性和可讀性。

如何改寫我的程式碼？

Eager execution 的使用方法對現有 TensorFlow 使用者來說應是直觀的。目前只有少量針對 eager 的 API；大多數現有的 API 和運算需要和啟用的 eager 一起工作。請記住以下內容：

一般對於 TensorFlow，我們建議如果你還沒有從排隊切換到使用 tf.data 進行輸入處理，請抓緊做。它更容易使用，也更快。檢視這篇博文（https://developers.googleblog.com/2017/09/introducing-tensorflow-datasets.html）和文件頁（https://www.tensorflow.org/programmers_guide/datasets）會有所幫助。

使用目標導向的層（比如 tf.layer.Conv2D() 或者 Keras 層），它們可以直接儲存變數。你可以為大多數模型寫程式碼，這對 eager execution 和圖構建同樣有效。也有一些例外，比如動態模型使用 Python 控制流改變基於輸入的計算。一旦呼叫 tfe.enable_eager_execution()，它不可被關掉。為了獲得圖行為，需要建立一個新的 Python session。

開始使用

這只是預釋出，還不完善。如果你想現在就開始使用，那麼：

  • 安裝 TensorFlow 的 nightly 版本（https://github.com/tensorflow/tensorflow#installation）
  • 檢視 README（包括 known issues），地址：https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/README.md
  • 從 eager execution 使用者指南（https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/g3doc/guide.md）中獲取詳細的指導
  • 在 GitHub 中檢視 eager 示例（https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples）
  • 及時檢視變更日誌（https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/README.md#changelog）檢視是否有更新

效能測試

Eager Execution 目前僅處於開發的前期，它的效能究竟如何？Google Brain 的工程師 Yaroslav Bulatov 對這一新工具做出了評測。TensorFlow 此前最令人詬病的問題就是它必須將計算定義為靜態圖。

我們在谷歌大腦的工作之一就是解決這類需求，並最終以命令式版本開源。但是這依賴於私有/不穩定的 API，而且這些 API 的維護成本會越來越高昂。幸運的是，PyTorch 滿足了研究員的需求，並且如今的 TensorFlow 也官方支援執行模式而不需要定義圖。

目前，Eager Execution 仍在積極開發中，但在最近釋出的可用版本非常有用，我們可以試用一下：

pip install tf-nightly-gpu
python
from tensorflow.contrib.eager.python import tfe
tfe.enable_eager_execution()
a = tf.random_uniform((10,))
b = tf.random_uniform((10,))
for i in range(100):
 a = a*a
 if a[0]>b[0]:
 break
print(i)

請注意，此操作並不需要處理圖，Session 就可以立即執行。若想應用 GPU 加速，請先將 tensor 拷貝至指定裝置。

a = a.gpu() # copies tensor to default GPU (GPU0)
a = a.gpu(0) # copies tensor to GPU0
a = a.gpu(1) # copies tensor to GPU1
a = a.cpu() # copies tensor back to CPU

埠命令程式碼

你可以將一個已有的 numpy/pytorch/matlab 的命令式程式碼重寫成正確的 API 呼叫。例如，

torch.sum -> tf.reduce_sum」
array.T -> tf.transpose(array) 等

我已使用 PyTorch 實現的 l-BFGS 作為練習，第一次在 GPU 上並行跑兩個實驗時（PyTorch & Eager），我得到前 8 位小數相同的結果。這使我大吃一驚，前所未聞。

終於！TensorFlow引入了動態圖機制Eager Execution

使用已有的基於圖的程式碼

如果你的程式碼不依賴於特定的 API，例如 graph_editor，你可以使用現有的程式碼並在 eager execution 模式下執行。

還有一個實驗性的函式「graph_callable」，可以將任意 tensorflow 子圖作為一個可以呼叫的函式。它仍然處於開發階段，但我能得到一個有效的例子來說明，該例子將 tensorflow /models 中的 resnet_model 包裝成一個 graph_callable。下面是一個隨機批大小訓練這個模型的例子。

一旦該功能上線，它應該有助於提高程式效能，具體可參考下文的效能部分。

擴充了梯度

原始 tf.gradients_function 的新衍生版本反映了autograd 的梯度。你可以呼叫在一個已有函式內呼叫「gradients_function」N 次獲得 N 階導數，即

# expensive way to compute factorial of n
def factorial(n):
 def f(x):
   return tf.pow(x, n)
 for i in range(n):
   f = tfe.gradients_function(f)
 return f(1.)

還有一個原始「custom_gradient」函式，這使得建立自定義梯度更容易。例如，假設我們想要平方函式，但在後向傳播時增加了噪聲。

@tfe.custom_gradient
def noisy_square(x):
  def grad(b):
      true_grad = 2*b*x
      return true_grad+tf.random_uniform(())
   return (x*x), grad
grad = tfe.gradients_function(noisy_square)
x = 2.
points = []
for i in range(20):
  x -= .9*grad(x)[0]
  print(x, loss(x))

效果如下：

終於！TensorFlow引入了動態圖機制Eager Execution

你會看到版本二收斂更慢，但是一旦收斂，它的泛化能力更好。

這種梯度修正對於實現如 KFAC 的高階優化演算法時十分有用。想想我早期所講，KFAC 在簡單網路中相當於啟用函式和反向傳播值白化的梯度下降。

這就可以理解為梯度在其兩邊乘上了白化的矩陣

終於！TensorFlow引入了動態圖機制Eager Execution

假設你已經將這些矩陣儲存為 m1，m2，那麼你自定義的乘操作可以是這樣的：

@tfe.custom_gradient
def kfac_matmul(W, A):
  def grad(B):
      true_grad1 = B @ tf.transpose(A)
      true_grad2 = tf.transpose(W) @ B
   return [m1 @ true_grad1 @ m2, true_grad2]
return W @ A, grad

注意，true_grad1, true_grad2 函式是乘法操作的反向傳播實現，請參考 Mike Giles 的第 4 頁「An extended collection of matrix derivative results for forward and reverse mode algorithmic differentiation」(https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf)

你可以通過使用 kfac_matmul 替代採用梯度下降演算法恢復原來的 kfac，或者你可以嘗試新的變種方法，利用動量和 Adam。

這裡（https://gist.github.com/yaroslavvb/eb02440272ddcbea549f1e47e4023376）有一個端到端的執行在 Eager execution 模式下的 KFAC 樣例。

效能

Eager Execution 模式使你的程式執行慢一點或慢很多的程度取決於你的計算高運算強度的卷積還是矩陣相乘。

做純矩陣乘法（超過 1 毫秒的時間）是沒有太大的差別，無論你用 tensorflow 快速模式，pytorch 或 tensorflow 經典模式。

終於！TensorFlow引入了動態圖機制Eager Execution

另一方面，端到端的例子更易受影響。

在測試中，當執行環境設定為 O(n^(1.5)) 操作，如 matmul/conv 時，Eager Execution 的速度要比 PyTorch 慢 20%，或者在大量 O(n) 操作如向量新增的例子中，比 PyTorch 慢 2-5 倍。

作為一個簡單的例子，我們使用吳恩達提出的 UFLDL 來訓練 MNIST 自編碼器。在批尺寸=60k，I-BFGS 的 history=5 時，大量的計算效能都被花在了自編碼器正向傳播上，Eager 的版本要比 PyTorch 慢 1.4 倍。

終於！TensorFlow引入了動態圖機制Eager Execution

在批尺寸為 60k，I-BFGS 的 history=100 的設定下，兩個迴環在每一步 I-BFGS（點積和向量增加）中執行「兩步遞迴」，Eager 版本的模型速度降低了 2.5 倍，而 PyTorch 僅受輕微影響。

終於！TensorFlow引入了動態圖機制Eager Execution

最後，如果我們將批尺寸減少到 10k，我們可以看到每次迭代的速度都要慢 5 倍，偶爾甚至會慢 10 倍，這可能是因為垃圾回收策略造成的。

終於！TensorFlow引入了動態圖機制Eager Execution

結論


雖然目前 Eager Execution 的表現還不夠強大，但這種執行模式可以讓原型設計變得容易很多。對於在 TensorFlow 中構建新計算任務的開發者而言，這種方式必將很快成為主流。

原文地址：

https://research.googleblog.com/2017/10/eager-execution-imperative-define-by.html

https://medium.com/@yaroslavvb/tensorflow-meets-pytorch-with-eager-mode-714cce161e6c


相關文章