標籤：

TensorFlow引入了動態圖機制Eager Execution

02-02

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。比如，要想使兩個矩陣相乘，我們這樣寫代碼：

使用 print 或者 Python 調試器檢查中間結果非常直接。

動態模型的構建可使用 Python 控制流。下面是使用 TensorFlow 算術操作的考拉茲猜想（Collatz conjecture）的一個示例：

這裡，tf.constant(12) 張量對象的使用將把所有數學運算提升為張量運算，從而所有的返回值將是張量。

梯度

多數 TensorFlow 用戶對自動微分（automatic differentiation）很感興趣。因為每次調用都有可能出現不同的運算，可以理解為我們把所有的正向運算錄到「磁帶」上，然後在計算梯度時進行「倒放」。梯度計算完成後，「磁帶」就沒用了。

如果你熟悉 autograd 包，我們提供的 API 與之非常類似。例如：

gradients_function 的調用使用一個 Python 函數 square() 作為參數，然後返回 Python callable，用於計算輸入的 square() 偏導數。因此，為了得到輸入為 3.0 時的 square() 導數，激活 grad(3.0)，也就是 6。

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

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

自定義梯度

用戶或許想為運算或函數自定義梯度。這可能有用，原因之一是它為一系列運算提供了更高效、數值更穩定的梯度。

下面的示例使用了自定義梯度。我們先來看函數 log(1 + e^x)，它通常用於計算交叉熵和 log 似然。

我們可以將自定義梯度應用於上述函數，簡化梯度表達式。注意下面的梯度函數實現重用了前向傳導中計算的 (tf.exp(x))，避免冗餘計算，從而提高梯度計算的效率。

建立模型

模型可以分成幾類。此處我們要提的模型可以通過創建一個簡單的兩層網路對標準的 MNIST 手寫數字進行分類。

我們推薦使用 tf.layers 中的類別（而非函數），這是因為它們創建並包含了模型參數（變數，variables）。變數的有效期和層對象的有效期緊密相關，因此需要對它們進行追蹤。

為什麼要使用 tfe.Network？一個網路包含了多個層，是 tf.layer.Layer 本身，允許將 Network 的對象嵌入到其它 Network 的對象中。它還包含能夠協助檢查、保存和修復的工具。

即使沒有訓練模型，我們也可以命令式地調用它並檢查輸出：

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

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

然後是訓練的循環過程：

implicit_gradients() 計算損失函數關於計算使用的所有 TensorFlow 變數的導數。

我們可以按往常使用 TensorFlow 的方式將計算轉移到 GPU 上：

（注意：我們簡化然後保存損失損失函數並直接調用 optimizer.minimize，但你也可以使用上面的 apply_gradients() 方法，它們是等價的。）

使用 Eager 和 Graphs

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

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

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

如何改寫我的代碼？

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

  • 一般對於 TensorFlow，我們建議如果你還沒有從排隊切換到使用 tf.data 進行輸入處理，請抓緊做。它更容易使用，也更快。查看這篇博文（developers.googleblog.com）和文檔頁（tensorflow.org/programm）會有所幫助。
  • 使用目標導向的層（比如 tf.layer.Conv2D() 或者 Keras 層），它們可以直接存儲變數。
  • 你可以為大多數模型寫代碼，這對 eager execution 和圖構建同樣有效。也有一些例外，比如動態模型使用 Python 控制流改變基於輸入的計算。
  • 一旦調用 tfe.enable_eager_execution()，它不可被關掉。為了獲得圖行為，需要建立一個新的 Python session。

開始使用

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

  • 安裝 TensorFlow 的 nightly 版本（github.com/tensorflow/t）
  • 查看 README（包括 known issues），地址：github.com/tensorflow/t
  • 從 eager execution 用戶指南（github.com/tensorflow/t）中獲取詳細的指導
  • 在 GitHub 中查看 eager 示例（github.com/tensorflow/t）
  • 及時查看變更日誌（github.com/tensorflow/t）查看是否有更新

性能測試

Eager Execution 目前僅處於開發的前期，它的性能究竟如何？Google Brain 的工程師 Yaroslav Bulatov 對這一新工具做出了評測。

TensorFlow 此前最令人詬病的問題就是它必須將計算定義為靜態圖。

我們在谷歌大腦的工作之一就是解決這類需求，並最終以命令式版本開源。但是這依賴於私有/不穩定的 API，而且這些 API 的維護成本會越來越高昂。

幸運的是，PyTorch 滿足了研究員的需求，並且如今的 TensorFlow 也官方支持執行模式而不需要定義圖。

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

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

埠命令代碼

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

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

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

使用已有的基於圖的代碼

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

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

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

拓展了梯度

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

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

效果如下：

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

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

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

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

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

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

這裡（gist.github.com/yarosla）有一個端到端的運行在 Eager execution 模式下的 KFAC 樣例。

性能

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

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

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

在測試中，當運行環境設置為 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 倍。

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

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

結論

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

原文地址：

  • research.googleblog.com
  • medium.com/@yaroslavvb/

選自Google Brain

作者：Asim Shankar & Wolff Dobson

機器之心編譯

本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。


推薦閱讀：

※在Docker中部署使用Tensorflow && Docker基本用法介紹
※TensorFlow 101
※tf.reset_default_graph

TAG:TensorFlow | PyTorch | GoogleBrain |