隨著機器學習等領(lǐng)域的飛速發(fā)展，亟需一門適用于該領(lǐng)域的編程語言。Julia官方博客發(fā)文探討了如何使用Julia重新思考ML工具，并對現(xiàn)代需要做的工作提供了一些見解。

自從Julia團隊提出“需要一流的語言、編譯器和機器學習（ML）生態(tài)系統(tǒng)”以來，該領(lǐng)域呈現(xiàn)出一些有趣的發(fā)展趨勢。

在現(xiàn)有的系統(tǒng)如TensorFlow或PyTorch中，不僅權(quán)衡問題（tradeoff）沒有得到解決，而且它們的“邊界”比以往任何時候都更加明顯，因為這兩個框架都包含不同的“靜態(tài)圖（static graph）”和“動態(tài)圖機制（eager execution）”接口。

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

Theano團隊的Myia將Python的一個子集區(qū)分并編譯為高性能GPU代碼；

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

Flux生態(tài)系統(tǒng)正在使用許多聚焦于ML的工具擴展Julia的編譯器，包括gradients、CUDA內(nèi)核編譯、自動批處理以及對TPU等新硬件的支持。

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

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

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

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

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

與其他下一代ML系統(tǒng)一樣，F(xiàn)lux致力于提供較為直觀的界面，并對任何類型的圖形構(gòu)建或性能注釋采取強硬措施。

Julia支持Flux所有特性，包括控制流、數(shù)據(jù)結(jié)構(gòu)和宏等。用戶可以在Jupyter notebook上交互式編程，并將高性能數(shù)字與便捷的繪圖、可視化做結(jié)合。

但Julia也想獲取傳統(tǒng)“靜態(tài)圖”框架的優(yōu)勢——零開銷的“源到源”AD、操作符融合、多GPU/分布式訓練和單二進制（single-binary）部署。

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

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

采用“梯度”

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

然而，這帶來了一個困難的權(quán)衡：我們要么接受解釋器的開銷（eager execution），要么凍結(jié)用戶控制流，并限制可以構(gòu)建的模型的種類（static graphs）。

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

將這個想法發(fā)揮到極致，我們構(gòu)建了Zygote，它直接在SSA形式的IR上工作，并支持控制流，遞歸，數(shù)據(jù)結(jié)構(gòu)和宏等語言功能。

然后，我們可以通過LLVM之類的編譯器生成的SSA形式的伴隨代碼，并將傳統(tǒng)編譯器優(yōu)化的所有好處應用于我們的前向和反向傳遞。

此外，這種方法為擴展該編譯器基礎(chǔ)結(jié)構(gòu)提供了機會，可以使用更高級和特定于域的優(yōu)化，例如內(nèi)核融合和編譯到TPU等加速器。 Swift for TensorFlow和Myia開發(fā)人員在源到源AD技術(shù)的復興中正在探索類似的方法。

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

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

為GPU編寫Julia

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

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

functionkernel_vadd(a,b,c)i=(blockIdx().x-1)*blockDim().x+threadIdx().xc[i]=a[i]+b[i]returnend

但是，Julia的類型特化（type specialization）可以在GPU上實現(xiàn)一組強大的附加抽象。例如，上面的代碼并不局限于密集的浮點數(shù)組，而是可以給出稀疏的復數(shù)數(shù)組。

Julia on TPUs

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

而我們只需要從已經(jīng)編寫的Julia程序中提取“靜態(tài)圖”并將其直接編譯到XLA，從而允許Julia本身在TPU上運行。（實際上，這只是Julia通常編譯過程的一個簡單擴展，該編譯過程從程序中提取盡可能大的“靜態(tài)子圖”，然后將它們發(fā)送到LLVM。）

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

自動Batching

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

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

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

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

總結(jié)

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