PyTorch 的動態圖一直是 TensorFlow 用戶夢寐以求的功能,谷歌也一直試圖在 TensorFlow 中實現相似的功能。最近,Google Brain 團隊發佈了 Eager Execution,一個由運行定義的新接口,讓 TensorFlow 開發變得簡單許多。在工具推出後,谷歌開發人員 Yaroslav Bulatov 對它的性能與 PyTorch 作了橫向對比。
今天,咱們爲 TensorFlow 引入了「Eager Execution」,它是一個命令式、由運行定義的接口,一旦從 Python 被調用,其操做當即被執行。這使得入門 TensorFlow 變的更簡單,也使研發更直觀。html
Eager Execution 的優勢以下:python
Eager Execution 如今處於試用階段,所以咱們但願獲得來自社區的反饋,指導咱們的方向。git
爲了更好地理解 Eager Execution,下面讓咱們看一些代碼。它很技術,熟悉 TensorFlow 會有所幫助。github
使用 Eager Execution算法
當你啓動 Eager Execution 時,運算會即刻執行,無需 Session.run() 就能夠把它們的值返回到 Python。好比,要想使兩個矩陣相乘,咱們這樣寫代碼:編程
使用 print 或者 Python 調試器檢查中間結果很是直接。性能優化
動態模型的構建可以使用 Python 控制流。下面是使用 TensorFlow 算術操做的考拉茲猜測(Collatz conjecture)的一個示例:網絡
這裏,tf.constant(12) 張量對象的使用將把全部數學運算提高爲張量運算,從而全部的返回值將是張量。session
梯度app
多數 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 中運行一樣的代碼。示例: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 一塊兒工做。請記住如下內容:
開始使用
這只是預發佈,還不完善。若是你想如今就開始使用,那麼:
性能測試
Eager Execution 目前僅處於開發的前期,它的性能究竟如何?Google Brain 的工程師 Yaroslav Bulatov 對這一新工具作出了評測。
TensorFlow 此前最使人詬病的問題就是它必須將計算定義爲靜態圖。
咱們在谷歌大腦的工做之一就是解決這類需求,並最終以命令式版本開源。可是這依賴於私有/不穩定的 API,並且這些 API 的維護成本會愈來愈高昂。
幸運的是,PyTorch 知足了研究員的需求,而且現在的 TensorFlow 也官方支持執行模式而不須要定義圖。
目前,Eager Execution 仍在積極開發中,但在最近發佈的可用版本很是有用,咱們能夠試用一下:
請注意,此操做並不須要處理圖,Session 就能夠當即執行。若想應用 GPU 加速,請先將 tensor 拷貝至指定設備。
端口命令代碼
你能夠將一個已有的 numpy/pytorch/matlab 的命令式代碼重寫成正確的 API 調用。例如,
我已使用 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」(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 經典模式。
另外一方面,端到端的例子更易受影響。
在測試中,當運行環境設置爲 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 中構建新計算任務的開發者而言,這種方式必將很快成爲主流。
原文地址:
選自Google Brain
做者:Asim Shankar & Wolff Dobson
機器之心編譯
本文爲機器之心編譯,轉載請聯繫本公衆號得到受權。