Julia官宣：為機器學習構建一種語言和編譯器

新智元報道

來源：julialang.org

編輯：三石

【新智元導讀】隨著機器學習等領域的飛速發展，亟需一門適用於該領域的編程語言。Julia官方博客發文探討了如何使用Julia重新思考ML工具，並對現代需要做的工作提供了一些見解。

自從Julia團隊提出「需要一流的語言、編譯器和機器學習（ML）生態系統」以來，該領域呈現出一些有趣的發展趨勢。

在現有的系統如TensorFlow或PyTorch中，不僅權衡問題（tradeoff）沒有得到解決，而且它們的「邊界」比以往任何時候都更加明顯，因為這兩個框架都包含不同的「靜態圖（static graph）」和「動態圖機制（eager execution）」介面。

在目前的框架還不夠完善的情況下，一些激動人心的新項目如雨後春筍般出現，完全省去了圖（graph），並將可微分編程（differentiable programming）引入主流：

Theano團隊的Myia將Python的一個子集區分並編譯為高性能GPU代碼；

Swift for TensorFlow擴展了Swift，可以將兼容的函數編譯為TensorFlow圖；

Flux生態系統正在使用許多聚焦於ML的工具擴展Julia的編譯器，包括gradients、CUDA內核編譯、自動批處理以及對TPU等新硬體的支持。

所有這些項目都擁有巨大的潛力，但團隊認為Julia更有優勢。

本文探討了團隊如何使用Julia重新思考ML工具，並對現代ML工具需要做的工作提供一些見解。

Flux加持，Julia更適於機器學習

我們需要一種語言來編寫可微演算法，而Flux使得Julia成為了這樣的一門語言。 Julia專為數學和數值計算而設計，非常適合表達ML演算法。同時，它在編譯器中融合了現代設計和新思想，更容易滿足最前沿ML的高性能需求。

在典型的框架中，所有的內容需要用幾十萬行的C++代碼來堆砌，而Flux僅僅是幾千行簡單的Julia代碼。只需要一個用於gradient的包（Zygote.jl），一個用於支持GPU的包（CuArrays.jl），「撒」上一些輕便的功能，「烘烤」十五分鐘，便可彈出一個功能齊全的ML堆棧。

與其他下一代ML系統一樣，Flux致力於提供較為直觀的界面，並對任何類型的圖形構建或性能注釋採取強硬措施。

Julia支持Flux所有特性，包括控制流、數據結構和宏等。用戶可以在Jupyter notebook上互動式編程，並將高性能數字與便捷的繪圖、可視化做結合。

但Julia也想獲取傳統「靜態圖」框架的優勢——零開銷的「源到源」AD、操作符融合、多GPU/分散式訓練和單二進位（single-binary ）部署。

這該如何實現呢？需要直接從Julia編寫的語法中提取和分析「靜態圖」，這實際上是編譯器完全正常的工作。從某些角度來看，大多數ML系統問題都是經過深入研究的標準編譯器問題。使用編譯語言就足以解決許多問題，擴展編譯器是解決更多問題的最佳方法。

在此只介紹這個領域當前工作中的一個示例—即獲取梯度、編譯GPU和TPU以及自動批處理。

採用「梯度」

我們突破了反向模式微分（reverse-mode differentiation）的極限，將其視為一個語言級別的問題。現有框架通過跟蹤（tracing）來實現這一點。引入了一種新的張量類型，它記錄了所執行的所有基本數學操作，產生了一個圖形（或符號表達式），其中刪除了主機語言的控制流和數據結構。

然而，這帶來了一個困難的權衡：我們要麼接受解釋器的開銷（eager execution），要麼凍結用戶控制流，並限制可以構建的模型的種類（static graphs）。

相反，如果圖（graph）是Julia自身的語法呢？

將這個想法發揮到極致，我們構建了Zygote，它直接在SSA形式的IR上工作，並支持控制流，遞歸，數據結構和宏等語言功能。

然後，我們可以通過LLVM之類的編譯器生成的SSA形式的伴隨代碼，並將傳統編譯器優化的所有好處應用於我們的前向和反向傳遞。

此外，這種方法為擴展該編譯器基礎結構提供了機會，可以使用更高級和特定於域的優化，例如內核融合和編譯到TPU等加速器。 Swift for TensorFlow和Myia開發人員在源到源AD技術的復興中正在探索類似的方法。

Julia用於此任務的一個關鍵優勢是它可用於實現基本數值庫，如微分方程求解器或優化庫; 這巧妙地解決了ML社區日益增長的需求，研究人員通過高性能代碼（如光線跟蹤器和物理引擎）反向傳播，但gradient仍必須在C ++中手動實現。

相比之下，由於Julia的實施是用Julia編寫的，所以從ODE到金融定價模型（ financial pricing model）的所有內容都可以輕鬆地進行區分。 將這些強大的工具帶入模型是深度學習真正成為可微分編程的地方。

為GPU編寫Julia

GPU編程是現代ML的重要組成部分。框架在內部提供內核，但是用戶只能看到有限的一組數學運算，不能直接對GPU進行編程。 相比之下，Julia中的GPU編程一直是一流的CUDA內核（可以很好地編寫並從腳本或筆記本中運行）。

一個簡單的向量加法核看起來與CUDA C等價。

function kernel_vadd(a, b, c)

i = (blockIdx().x-1) * blockDim().x + threadIdx().x

c[i] = a[i] + b[i]

return

end

但是，Julia的類型特化（type specialization）可以在GPU上實現一組強大的附加抽象。例如，上面的代碼並不局限於密集的浮點數組，而是可以給出稀疏的複數數組。

Julia on TPUs

谷歌最近開放了他們的雲TPU使用的XLA IR，使得ML以外的其他框架和用戶都可以利用這個重量級硬體。 XLA功能強大但有限：它無法運行Python解釋器，當然也沒有良好的性能。 然後框架最終處於與gradient相似的位置，只能使用程序跟蹤來撬開Python，最終得到一個快速但更有限的ML語言。

而我們只需要從已經編寫的Julia程序中提取「靜態圖」並將其直接編譯到XLA，從而允許Julia本身在TPU上運行。（實際上，這只是Julia通常編譯過程的一個簡單擴展，該編譯過程從程序中提取儘可能大的「靜態子圖」，然後將它們發送到LLVM。）

這使我們可以充分利用Julia語言的表現力，包括 控制流，遞歸，多調度，高階函數，強大的數據結構和抽象，自定義數字類型，以及現有的包，如微分方程求解器和線性代數常式。所有這些工作都是在TPU中使用高性能收縮陣列引擎的同時進行的。

自動Batching

為了從這些加速器中獲得最大收益，批處理程序通常會同時將前向和反向傳遞應用於多個訓練示例。在一些簡單的情況下，比如卷積網路，通過將10張圖像按照額外的批處理維度連接起來來處理這個問題是很簡單的。但是在處理各種結構的輸入（如樹或圖）時，這項任務變得更加困難。

大多數研究人員通過手工批處理代碼來解決這一問題。針對不同的框架（DyNet、TensorFlow Fold）提出了不同的解決方案，它們在可能的情況下嘗試將一些高級操作批處理在一起，但是這些操作通常要麼有自己的可用性問題，要麼無法實現手工編寫的代碼的性能。

我們認為這個問題與單程序多數據（SPMD）編程的問題是相同的，後者已經被語言和編譯器社區研究了幾十年，並且在最近的批處理方法（如matchbox）中變得很明顯。 實際上，它與GPU內部使用的並行模型非常相似，並且已經實現為CPU的SIMD單元的編譯器變換。

從這項工作中獲得靈感，我們正在Julia中實現相同的轉換，為標量SIMD單元和模型級批處理提供SPMD編程。這使我們能夠實現在單個示例上編寫簡單代碼的理想，同時仍然在現代硬體上獲得最佳性能。

總結

我們認為，機器學習的未來取決於語言和編譯技術，特別是擴展新的或現有的語言，以滿足ML研究的高要求。這不僅對ML社區有好處，對一般的數值編程也有好處；能夠很好地支持差異化、向量化和外來硬體的語言將足以推動科學的許多進步。

https://julialang.org/blog/2018/12/ml-language-compiler

【加入社群】

新智元 AI 技術 + 產業社群招募中，歡迎對 AI 技術 + 產業落地感興趣的同學，加小助手微信號：aiera2015_2入群；通過審核後我們將邀請進群，加入社群後務必修改群備註（姓名 - 公司 - 職位；專業群審核較嚴，敬請諒解）。

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