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

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

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

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

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

Swift for TensorFlow擴(kuò)展了Swift，可以將兼容的函數(shù)編譯為T(mén)ensorFlow圖；

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

所有這些項(xiàng)目都擁有巨大的潛力，但團(tuán)隊(duì)認(rèn)為Julia更有優(yōu)勢(shì)。

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

Flux加持，Julia更適于機(jī)器學(xué)習(xí)

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

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

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

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

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

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

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

采用“梯度”

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

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

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

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

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

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

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

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

為GPU編寫(xiě)Julia

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

一個(gè)簡(jiǎn)單的向量加法核看起來(lái)與CUDA C等價(jià)。

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

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

Julia on TPUs

谷歌最近開(kāi)放了他們的云TPU使用的XLA IR，使得ML以外的其他框架和用戶(hù)都可以利用這個(gè)重量級(jí)硬件。 XLA功能強(qiáng)大但有限：它無(wú)法運(yùn)行Python解釋器，當(dāng)然也沒(méi)有良好的性能。 然后框架最終處于與gradient相似的位置，只能使用程序跟蹤來(lái)撬開(kāi)Python，最終得到一個(gè)快速但更有限的ML語(yǔ)言。

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

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

自動(dòng)Batching

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

大多數(shù)研究人員通過(guò)手工批處理代碼來(lái)解決這一問(wèn)題。針對(duì)不同的框架（DyNet、TensorFlow Fold）提出了不同的解決方案，它們?cè)诳赡艿那闆r下嘗試將一些高級(jí)操作批處理在一起，但是這些操作通常要么有自己的可用性問(wèn)題，要么無(wú)法實(shí)現(xiàn)手工編寫(xiě)的代碼的性能。

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

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

總結(jié)

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