有了Julia語言，深度學習框架從此不需要計算圖

選自julialang

作者：Mike Innes 等

機器之心編譯

參與：劉曉坤、思源

本文基於 NeurIPS MLSys 的一篇論文《Fashionable Modelling with Flux》，探討開發者們如何使用Julia語言從頭開始思考機器學習工具，並提供對於現代機器學習工具所需改進的一些見解，涉及新的可微分編程工具 Flux、求梯度、支持 GPU 和 TPU、自動批處理。為什麼 Julia 式的機器學習不需要計算圖呢？因為 Julia 的機器學習語法就是計算圖。

鑒於機器學習（ML）對編程語言、編譯器和生態系統的眾多需求，現在已經有很多有趣的發展。不僅 TensorFlow 和 PyTorch 等現有系統間的權衡得不到解決，而且這兩個框架都包含不同的「靜態圖」和「eager execution」介面，但它們的形式已經比以前更加清晰。與此同時，機器學習模型基本上是可微分演算法的思想（通常稱為可微分編程）已經流行起來。

當前的機器學習框架遇到了阻礙，很多已有的新項目都完全移除了計算圖，從而使可微分編程成為主流。例如，由 Theano 團隊開發的 Myia 可以求微分並編譯 Python 的一個子集為高性能 GPU 代碼。Swift for TensorFlow 作為 Swift 語言的擴展，它可以將兼容的函數編譯為 TensorFlow 計算圖。最後，Flux 生態系統為 Julia 編譯器提供了一些機器學習專用的工具，包括：first-class gradients、即時CUDA核編譯、自動批處理（automatic batching）以及對新硬體（例如 TPU）的支持。

所有這些項目都有巨大的潛力，但目前看來 Julia 具有優勢。

Flux 簡介

我們需要一種語言來編寫可微分演算法，Flux 使 Julia 變成了這樣的語言。Julia 專為數學和數值計算而設計，非常適合表達機器學習演算法。同時，它在編譯器中融合了現代設計和新思想，可以更輕鬆地滿足尖端 ML 的高性能需求。

典型的框架通常包含數十萬行 C++代碼，Flux 卻只有千行 Julia 代碼。只需要一個求梯度的包（Zygote.jl）、一個用於 GPU 支持的包（CuArrays.jl）、再加上一些輕量函數，我們就能得到一個功能齊全的機器學習堆棧。

與其他下一代機器學習系統一樣，Flux 致力於提供直觀（「eager」或「define-by-run」）的介面，並對任何類型的計算圖構建或性能注釋進行嚴格控制。從控制流、數據結構到宏，Flux 支持語言的所有特徵。用戶可以在 Jupyter 筆記本中互動式地寫代碼，並將高性能數值計算與方便的繪圖、可視化相結合。但我們也希望獲得傳統上由「靜態圖」框架所帶來的好處，例如零開銷源到源 AD、OP 融合、多 GPU /分散式訓練和二進位部署等。

我們怎麼能做到這一切？實際上，我們需要直接從 Julia 語法中提取和分析「靜態圖」，這實際完全上是編譯器的正常工作。通過適當的角度來看，大多數機器學習系統問題都是標準的且經過充分研究的編譯器問題。使用編譯語言足以解決許多問題，擴展該編譯器是解決更多問題的最佳方法。本文僅介紹了我們目前在該領域的工作範例，即求梯度、為 GPU 和 TPU 提供代碼編譯，以及自動批處理。

求梯度

推動反向模式求微分的極限，我們將此視為語言層面的問題。求微分是一種符號轉換，屬於編譯器的領域。現有框架通過追蹤（實際上是一種部分評估或抽象解釋）來實現這一目標。人們引入了一種新的張量類型，它記錄了所執行的所有基本數學運算，生成一個計算圖（或符號表達式），其中刪除了宿主語言的控制流和數據結構。然而，這給出了一個艱難的權衡：我們要麼接受解釋器的開銷（eager execution），要麼固定用戶的控制流並限制可以構建的模型種類（靜態圖）。

反之，如果「計算圖」就是 Julia 自己的語法呢？通過將這個想法發揮到極致，我們構建了 Zygote，它直接在 SSA 形式的中間表徵（IR）上工作，支持控制流、遞歸、數據結構和宏等語言功能。然後，我們可以通過 LLVM 之類的編譯器生成 SSA 形式的伴隨代碼，並將傳統編譯器優化的所有優勢應用於前向和後向傳播。此外，這種方法還為擴展該編譯器基礎結構提供了可能，可以使用更高級和特定領域的優化，例如用於 TPU 等加速器的內核融合和編譯。TensorFlow 的 Swift 和 Myia 開發人員在源到源 AD 技術的復興中正在探索類似的方法。

Julia 用於此任務的一個關鍵優勢是它可用於實現基本數值計算庫，如微分方程求解器或優化庫；這巧妙地解決了機器學習社區不斷增長的需求，研究人員通過高性能代碼（如光線追蹤和物理引擎）進行反向傳播，但求梯度仍必須在 C++中手動實現。相比之下，由於 Julia 的實現是用 Julia 編寫的，因此可以輕鬆對從 ODE 到金融定價模型等求微分。將這些強大的工具帶入模型是深度學習真正成為可微分編程的關鍵。

編譯 Julia 到 GPU 上

GPU 編程是現代機器學習的重要組成部分，但 GPU 通常被視為實現細節。因為框架在內部提供內核，但用戶只能使用一組有限的數學運算，無法直接對 GPU 進行編程。相比之下，Julia 中的 GPU 編程一直是一流的 CUDA 內核（可以很好地編寫並從腳本或 notebook 中運行）。如下簡單的向量加法內核看起來類似於 CUDA C：

但是，Julia 的類型特化可以在 GPU 上實現一組強大的附加抽象。例如，上面的代碼不限於浮點數的密集數組，而是可以給出複數的稀疏數組；Julia 的常規特化機制將動態地生成一組新的 PTX 指令。我們甚至可以將此代碼進一步抽象為可利用「+」函數的「高階內核」，從而在四行代碼內創建一整套函數 map（f，x，y）。

這可以實現一些強大的技巧，即使你自己從不編寫 CUDA 代碼。例如，我們可以透明地將大型廣播（broadcast）表達式（例如 1 /（1 + exp（-x））及其向後傳遞融合到單個 GPU 內核中，從而獲得顯著加速。我們期望原生 GPU 代碼生成能力和生態系統將為各種基於 Julia 的機器學習庫提供支持。

編譯 Julia 到 TPU 上

更進一步，谷歌最近開放了雲 TPU 使用的 XLA IR，使得其他框架和用戶都可以利用這個重量級硬體。XLA 功能強大但有限制：它無法運行 Python 解釋器，當然也沒有良好的性能。

而我們只需要從編寫的 Julia 程序中提取「靜態圖」並將其直接編譯為 XLA，從而允許 Julia 本身在 TPU 上運行。（事實上，這只是 Julia 一般編譯過程的簡單擴展，它在將程序發送到 LLVM 之前從程序中提取最大的「靜態子圖」。）這使我們可以充分利用 Julia 語言的表現力，包括控制流、遞歸、多調度、高階函數、強大的數據結構和抽象、自定義數值類型，以及現有的包，如微分方程求解器和線性代數常式。所有這些都在獲得高性能收縮陣列引擎的優勢的同時，在 TPU 內運行。你今天就可以嘗試，其中包括 ResNet 等大型機器學習模型和 TSVD 等線性代數常式。

項目地址：https://github.com/JuliaTPU/XLA.jl

自動批處理（Automatic Batching）

為了從這些加速器中獲得最大收益（每個內核啟動可能會產生大量開銷，但是在輸入大小上可以很好地擴展），批處理程序通常會同時將前向和反向傳播應用於多個訓練樣本。在簡單的情況下，例如使用卷積網路，通過在額外的批量維度上拼接 10 張圖像來處理這個問題會變得很簡單。但是，當處理可變結構的輸入（例如樹或圖形）時，此任務變得更加困難。

大多數研究人員通過人工完成批處理代碼來解決這個問題，這樣做的成本非常大。人們已經針對不同的框架提出了不同的解決方案（DyNet、TensorFlow Fold，它試圖在可能的情況下將一些高級 OP 一起批處理，但是這些通常要麼具有其自身的可用性問題，要麼沒有實現手寫代碼的性能。

我們認為這個問題與單程序多數據（SPMD）編程的問題完全相同，單程序多數據編程幾十年來一直被語言和編譯器社區充分研究。實際上，它與 GPU 內部使用的並行模型非常相似，並且已經實現 CPU 的 SIMD 單元的編譯器變換。通過從這項工作中汲取靈感，我們在 Julia 中實現了相同的變換，為標量 SIMD 單元和模型級批處理提供 SPMD 編程。這使我們能夠編寫對單個樣本進行操作的簡單代碼，同時仍然在現代硬體上獲得最佳性能。

結論

我們相信機器學習的未來取決於編程語言和編譯器技術，尤其是擴展新的或現有的語言以滿足機器學習研究的高要求。這不僅適用於機器學習社區，也適用於一般的數值規劃；能夠支持微分、向量化和新型硬體的編程語言將足以推動科學的許多進步。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！