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社群有好處，對一般的數值程式設計也有好處；能夠很好地支援差異化、向量化和外來硬體的語言將足以推動科學的許多進步。
原文連結：