?Julia從一出生開始，就瞄準(zhǔn)了科學(xué)計(jì)算領(lǐng)域，并且一直在與Python暗中較量。

在神經(jīng)網(wǎng)絡(luò)的框架上，Python有PyTorch和TensorFlow，幾乎是深度學(xué)習(xí)開發(fā)的首選框架，并且獲得了Meta和Google在技術(shù)和資金上的支持，蓬勃發(fā)展。

雖然Julia也有Flux.jl框架，但Julia社區(qū)一直依賴于語言本身的高性能產(chǎn)生的生產(chǎn)力，所以Flux.jl的代碼量相比Python框架來說，可以稱得上是特別「苗條」了，例如PyTorch和TensorFlow包括了整個獨(dú)立的語言和編譯器（torchscript、XLA等），而Flux.jl僅僅由Julia語言編寫。

當(dāng)然，世界上沒有免費(fèi)的午餐，如果以不同的視角來看，想要在機(jī)器學(xué)習(xí)領(lǐng)域開發(fā)出一個簡單、通用且高性能的框架幾乎是不可能的，只能不斷權(quán)衡。

比如對于一個特定的問題，如果需要稀疏的小模型，想要獲得最高性能的方法就是重寫一遍，而非采用通用框架。

最近Julia社區(qū)又開源了一個新框架SimpleChains.jl，在小模型場景下相比PyTorch最少能提速5倍。

代碼鏈接：https://github.com/PumasAI/SimpleChains.jl

開發(fā)人員表示，這個框架不會對所有人都有用，但對那些需要它的人來說，它是非常有用的。

有網(wǎng)友表示十分贊同：「不同的任務(wù)用不同的工具」，因?yàn)門F和pyTorch消耗了大量的內(nèi)存，并且沒有原地操作，所以在小模型上很浪費(fèi)時間。他數(shù)年前在Netflex時就設(shè)計(jì)開發(fā)了一個D語言的框架vectorflow，目前在github已獲1200個stars

機(jī)器學(xué)習(xí)模型的假設(shè)

SimpleChains.jl是由Pumas-AI和Julia Computing與Roche和馬里蘭大學(xué)巴爾的摩分校合作開發(fā)的一個庫，它的主要目的就是為小型神經(jīng)網(wǎng)絡(luò)提供盡可能高的性能。

SimpleChains.jl最開始用于在醫(yī)療數(shù)據(jù)分析中用于科學(xué)機(jī)器學(xué)習(xí)（SciML）的解決方案：小型神經(jīng)網(wǎng)絡(luò)（和其他近似器，如傅里葉數(shù)列或切比雪夫多項(xiàng)式展開）可以與已知的半生理學(xué)模型（semi-physiologic models）相結(jié)合，發(fā)現(xiàn)以前未知的機(jī)制和預(yù)后因素。

從黑洞動力學(xué)到地震安全建筑的開發(fā)，SciML方法的有效性已經(jīng)在許多學(xué)科中得到證實(shí)，能夠靈活地發(fā)現(xiàn)/指導(dǎo)（生物）物理方程。

應(yīng)用場景變化太大，在這種情況下，使用一些專用（specialization）的神經(jīng)網(wǎng)絡(luò)才有可能提升模型的運(yùn)行性能。

具體來說，在機(jī)器學(xué)習(xí)模型的研究中，通常依賴于一個假設(shè)：神經(jīng)網(wǎng)絡(luò)足夠大，其中矩陣乘法（如卷積）的O(n^3)時間成本占了運(yùn)行時間的絕大部分，這基本上也是機(jī)器學(xué)習(xí)庫的大部分機(jī)制背后的4大指導(dǎo)原則：

1. 矩陣乘法的復(fù)雜度是立方的，而內(nèi)存分配的規(guī)模是線性的，所以用非分配（non-allocating）內(nèi)存的方式來操作向量的優(yōu)先級并不高；

2. 目前AI加速的工作主要集中于GPU內(nèi)核加速，讓指令運(yùn)行盡可能快，由于這些大型矩陣-矩陣操作在GPU上是最快的，并且也是大模型的主要瓶頸，所以性能基準(zhǔn)基本上只是衡量這些特定內(nèi)核的速度；

3. 當(dāng)做自動微分反向傳播時，將數(shù)值復(fù)制到內(nèi)存的操作幾乎感覺不到，內(nèi)存分配被較大的內(nèi)核調(diào)用所隱藏；

4. 用戶可以隨意寫一個tape來生成反向傳播，雖然增加了在前向過程中建立字典的成本，但是也會被更大的內(nèi)核調(diào)用所掩蓋。

但，這些假設(shè)在真實(shí)的案例中是否真的能全部成立？

如果不成立的話，能不能把重點(diǎn)放在這些方面的改進(jìn)，從而提升更高的運(yùn)算性能？

小型神經(jīng)網(wǎng)絡(luò)的瓶頸在哪？

對于初學(xué)者來說，可以先測試一下假設(shè)1和2，通過一段Julia代碼來測試內(nèi)存申請時間、GPU運(yùn)算時間等。

可以看到當(dāng)我們進(jìn)行較大的矩陣乘法操作時，比如100x100*100x100，基本可以忽略由于內(nèi)存分配而產(chǎn)生的任何開銷。

但同時也可以看到，在lower end有可能出現(xiàn)一些相當(dāng)顯著的性能提升，這些收益是通過使用純Julia LoopVectorization.jl實(shí)現(xiàn)的，因?yàn)闃?biāo)準(zhǔn)的BLAS工具在這個區(qū)域往往有額外的線程開銷（同樣，在這個區(qū)域沒有進(jìn)行優(yōu)化）。

如果你一直在利用GPU帶來的好處而不去研究細(xì)節(jié)，那么這個事實(shí)可能會讓你大吃一驚！GPU被設(shè)計(jì)成具有許多內(nèi)核的慢速芯片，因此它們只對非常并行的操作有效，例如大型矩陣乘法。正是從這一點(diǎn)出發(fā)，假設(shè)2可以被認(rèn)為是大型網(wǎng)絡(luò)操作。

但同樣，在小網(wǎng)絡(luò)的情況下，由于缺乏并行計(jì)算，使用GPU內(nèi)核的性能可能還不如設(shè)計(jì)良好的CPU內(nèi)核。

矩陣操作只有在能夠使用批處理（A*B中的B矩陣的每一列都是一個單獨(dú)的批處理）時才會發(fā)生。

在大部分科學(xué)機(jī)器學(xué)習(xí)的情境下，如ODE鄰接中的向量Jacobian乘積的計(jì)算，這種操作是矩陣-向量乘法。這些操作的時間復(fù)雜度只有O(n^2)，在這種情況下內(nèi)存開銷會被放大。

神經(jīng)網(wǎng)絡(luò)的基本操作是Sigma，所以還有一個O(n)時間復(fù)雜度的操作，這種情況下內(nèi)存開銷顯得更嚴(yán)重。

對于假設(shè)3和4來說，需要更加關(guān)注反向傳播的實(shí)現(xiàn)。不同機(jī)器學(xué)習(xí)庫的自動微分方法也存在著區(qū)別。有些庫是立刻反向傳播梯度值，也有些需要把梯度保存起來，這樣就又需要額外的內(nèi)存開銷操作了。

在特定的應(yīng)用里面，如果知道梯度立刻傳播，就可以立即計(jì)算梯度，相比通用實(shí)現(xiàn)來說，只需要一個緩存向量解千愁，原地賦值，這樣的話所有自動微分的額外開銷都沒有了。

基于這些想法，研究人員開源了SimpleChains.jl，可以很好地解決這類優(yōu)化問題，可以在CPU上快速擬合和優(yōu)化小模型，早期的神經(jīng)網(wǎng)絡(luò)原型模型設(shè)計(jì)大多都希望：

1. 達(dá)到更好的性能，最好能達(dá)到CPU的峰值FLOPs；

2. 專注于小尺寸的模型，在早期開發(fā)階段放棄一些針對大型模型的內(nèi)核優(yōu)化操作（如緩存平鋪）；

3. 有一個API，其中的向量的參數(shù)和梯度都是first class，以便更容易地與各種優(yōu)化器或求解器（如BFGS）協(xié)同工作；

4. 使用「純Julia」編寫，更方便開發(fā)和優(yōu)化；在大量使用LoopVectorization.jl的同時，SimpleChains.jl并不依賴任何BLAS或NN庫。

開發(fā)人員的長期目標(biāo)是將這種循環(huán)編譯器的優(yōu)化方法擴(kuò)展到自動產(chǎn)生pullbacks。但這種以編譯器為中心的方法已經(jīng)被用于實(shí)現(xiàn)的便利性：雖然我們?nèi)匀恍枰謱懱荻?，但我們不需要對它們進(jìn)行手工優(yōu)化。

SimpleChains.jl的實(shí)際性能怎么樣？

研究人員用一個2×2的矩陣做了一個實(shí)驗(yàn)，在帶有AVX512指令集的Intel i9-10980XE跑了一下，1萬個epoch花了0.41秒，相比之下pyTorch花了15秒，也就是說在這種微型神經(jīng)網(wǎng)絡(luò)上，提速大約35倍。

把實(shí)驗(yàn)換到AMD EPYC 7513 帶有AVX2指令的機(jī)器上，Julia的實(shí)現(xiàn)花費(fèi)時間為0.72秒，而PyTorch的實(shí)現(xiàn)則需要70秒，差距拉升到了100倍。

研究人員又在AMD Ryzen 9 5950X實(shí)驗(yàn)了一份Jax代碼，Julia耗時為1.3秒，Jax則需要14秒，提升約10倍。

換到Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz 平臺上，Jax耗時為9秒，Julia需要0.4秒，大約22倍提升。

再換到差一點(diǎn)的處理器，6核CPU上，Jax需要19秒，而Julia需要9秒，速度提升就只有2倍了。

在稍微大一點(diǎn)的、實(shí)際可用的神經(jīng)網(wǎng)絡(luò)上，訓(xùn)練速度還會有這么大的差距嗎？

研究人員用LeNet5來測試MNIST，這個例子只是一個非常保守的速度估計(jì)，因?yàn)樵诟鼈鹘y(tǒng)的機(jī)器學(xué)習(xí)用例中，批處理可以使用矩陣乘法，不過即使在這種情況下，由于semi-small的網(wǎng)絡(luò)規(guī)模，也能看到大量的性能優(yōu)勢。

在batch size為2048的情況下訓(xùn)練10個epoch，用PyTorch在A100上訓(xùn)練兩次耗時為17.66和17.62，準(zhǔn)確率分別為94.91%和96.92%；在V100上訓(xùn)練時間為16.29和15.94，準(zhǔn)確率分別為95.6%和97.5%

不過這個問題對于GPU來說還是殺雞用牛刀了，在2048的batch size上運(yùn)算速度還是很快，時間主要耗費(fèi)在CPU轉(zhuǎn)移到GPU上了。

在AMD EPYC 7513和Intel i9 10980XE又進(jìn)行了兩次實(shí)驗(yàn)，結(jié)果比GPU更快，準(zhǔn)確率也更高。

換到SimpleChains.jl，在AMD平臺上耗時為3秒，準(zhǔn)確率98.3%；在Intel平臺上，耗時僅為1秒，準(zhǔn)確率為98.2%；即使在筆記本的Intel平臺上，耗時也僅為5.3秒，準(zhǔn)確率97%

目前大型機(jī)器學(xué)習(xí)框架在專注于為其99.9%的用戶提供一流的性能方面做得非常好，但在另外0.1%的小模型用戶手里，框架卻不好用。

這就是可組合性和靈活性的優(yōu)勢：一種允許你輕松構(gòu)建機(jī)器學(xué)習(xí)框架的語言，也是一種允許你構(gòu)建替代框架的語言，這些框架針對替代人群進(jìn)行優(yōu)化。?