最近超級火爆的DALL·E Mini都知道吧，它的模型就是基于JAX進(jìn)行編程的，從而充分地利用了谷歌TPU帶來的優(yōu)勢。

TensorFlow的黃昏和PyTorch的崛起

2015年，谷歌開發(fā)的機(jī)器學(xué)習(xí)框架——TensorFlow問世。

當(dāng)時，TensorFlow只是Google Brain的一個小項目。

誰也沒有想到，剛一問世，TensorFlow就變得非?；鸨?。

優(yōu)步、愛彼迎這種大公司在用，NASA這種國家機(jī)構(gòu)也在用。而且還都是用在他們各自最為復(fù)雜的項目上。

而截止到2020年11月，TensorFlow的下載次數(shù)已經(jīng)達(dá)到了1.6億次。

不過，谷歌好像并沒有十分在乎這么多用戶的感受。

奇奇怪怪的界面和頻繁的更新都讓TensorFlow對用戶越來越不友好，并且越來越難以操作。

甚至，就連谷歌內(nèi)部，也覺得這個框架在走下坡路。

其實谷歌如此頻繁的更新也實屬無奈，畢竟只有這樣才能追得上機(jī)器學(xué)習(xí)領(lǐng)域快速地迭代。

于是，越來越多的人加入了這個項目，導(dǎo)致整個團(tuán)隊慢慢失去了重點(diǎn)。

而原本讓TensorFlow成為首選工具的那些閃光點(diǎn)，也被埋沒在了茫茫多的要素里，不再受人重視。

這種現(xiàn)象被Insider形容為一種「貓鼠游戲」。公司就像是一只貓，不斷迭代出現(xiàn)的新需求就像是一只只老鼠。貓要時刻保持警惕，隨時撲向老鼠。

這種困局對最先打入某一市場的公司來說是避不開的。

舉個例子，就搜索引擎來說，谷歌并不是第一家。所以谷歌能夠從前輩（AltaVista、Yahoo等等）的失敗中總結(jié)經(jīng)驗，應(yīng)用在自身的發(fā)展上。

可惜到了TensorFlow這里，谷歌是被困住的那一個。

正是因為上面這些原因，原先給谷歌賣命的開發(fā)者，慢慢對老東家失去了信心。

昔日無處不在的TensorFlow漸漸隕落，敗給了Meta的后起之秀——PyTorch。

2017年，PyTorch的測試版開源。

2018年，F(xiàn)acebook的人工智能研究實驗室發(fā)布了PyTorch的完整版本。

值得一提的是，PyTorch和TensorFlow都是基于Python開發(fā)的，而Meta則更注重維護(hù)開源社區(qū)，甚至不惜大量投入資源。

而且，Meta關(guān)注到了谷歌的問題所在，認(rèn)為不能重蹈覆轍。他們專注于一小部分功能，并把這些功能做到最好。

Meta并沒有步谷歌的后塵。這款首先在Facebook開發(fā)出來的框架，慢慢成為了行業(yè)標(biāo)桿。

一家機(jī)器學(xué)習(xí)初創(chuàng)公司的研究工程師表示，「我們基本都用PyTorch。它的社群和開源做得是最出色的。不僅有問必答，給的例子也很實用。」

面對這種局面，谷歌的開發(fā)者、硬件專家、云提供商，以及任何和谷歌機(jī)器學(xué)習(xí)相關(guān)的人員在接受采訪時都說了一樣的話，他們認(rèn)為TensorFlow失掉了開發(fā)者的心。

經(jīng)歷了一系列的明爭暗斗，Meta最終占了上風(fēng)。

有專家表示，谷歌未來繼續(xù)引領(lǐng)機(jī)器學(xué)習(xí)的機(jī)會正慢慢流失。

PyTorch逐漸成為了尋常開發(fā)者和研究人員的首選工具。

從Stack Overflow提供的互動數(shù)據(jù)上看，在開發(fā)者論壇上有關(guān)PyTorch的提問越來越多，而關(guān)于TensorFlow的最近幾年一直處于停滯狀態(tài)。

就連文章開始提到的優(yōu)步等等公司也轉(zhuǎn)向PyTorch了。

甚至，PyTorch后來的每一次更新，都像是在打TensorFlow的臉。

谷歌機(jī)器學(xué)習(xí)的未來——JAX

就在TensorFlow和PyTorch打得熱火朝天的時候，谷歌內(nèi)部的一個「小型黑馬研究團(tuán)隊」開始致力于開發(fā)一個全新的框架，可以更加便捷地利用TPU。

2018年，一篇題為《Compiling machine learning programs via high-level tracing》的論文，讓JAX項目浮出水面，作者是Roy Frostig、Matthew James Johnson和Chris Leary。

從左至右依次是這三位大神

而后，PyTorch原始作者之一的Adam Paszke，也在2020年初全職加入了JAX團(tuán)隊。

JAX提供了一個更直接的方法用于處理機(jī)器學(xué)習(xí)中最復(fù)雜的問題之一：多核處理器調(diào)度問題。

如此帶來的好處就是，讓盡可能多的TPU片刻間就能得到響應(yīng)，從而燃燒我們的「煉丹小宇宙」。

最終，相比于臃腫的TensorFlow，JAX解決了谷歌內(nèi)部的一個心頭大患：如何快速訪問TPU。

下面簡單介紹一下構(gòu)成JAX的Autograd和XLA。

Autograd主要應(yīng)用于基于梯度的優(yōu)化，可以自動區(qū)分Python和Numpy代碼。

它既可以用來處理Python的一個子集，包括循環(huán)、遞歸和閉包，也可以對導(dǎo)數(shù)的導(dǎo)數(shù)進(jìn)行求導(dǎo)。

此外，Autograd支持梯度的反向傳播，這也就這意味著它可以有效地獲取標(biāo)量值函數(shù)相對于數(shù)組值參數(shù)的梯度，以及前向模式微分，并且兩者可以任意組合。

XLA（Accelerated Linear Algebra）可以加速TensorFlow模型而無需更改源代碼。

當(dāng)一個程序運(yùn)行時，所有的操作都由執(zhí)行器單獨(dú)執(zhí)行。每個操作都有一個預(yù)編譯的GPU內(nèi)核實現(xiàn)，執(zhí)行器會分派到該內(nèi)核實現(xiàn)。

舉個栗子：

def model_fn(x, y, z):  return tf.reduce_sum(x + y * z)

在沒有XLA的情況下運(yùn)行，該部分會啟動三個內(nèi)核：一個用于乘法，一個用于加法，一個用于減法。

而XLA可以通過將加法、乘法和減法「融合」到單個GPU內(nèi)核中，從而實現(xiàn)優(yōu)化。

這種融合操作不會將由內(nèi)存產(chǎn)生的中間值寫入y*z內(nèi)存x+y*z；相反，它將這些中間計算的結(jié)果直接「流式傳輸」給用戶，同時將它們完全保存在GPU中。

在實踐中，XLA可以實現(xiàn)約7倍的性能改進(jìn)和約5倍的batch大小改進(jìn)。

2020年DeepMind正式投入JAX的懷抱，而這也宣告了谷歌親自下場，自此之后各種開源的庫層出不窮。

谷歌的挑戰(zhàn)

作為一個新秀，Jax雖然可以借鑒PyTorch和TensorFlow這兩位老前輩的優(yōu)點(diǎn)，但有的時候后發(fā)可能也會帶來劣勢。

首先，JAX還太「年輕」，作為實驗性的框架，遠(yuǎn)沒有達(dá)到一個成熟的谷歌產(chǎn)品的標(biāo)準(zhǔn)。

除了各種隱藏的bug以外，JAX在一些問題上仍然要依賴于其他框架。

拿加載和預(yù)處理數(shù)據(jù)來說，就需要用TensorFlow或PyTorch來處理大部分的設(shè)置。

顯然，這和理想的「一站式」框架還相去甚遠(yuǎn)。

其次，JAX主要針對TPU進(jìn)行了高度的優(yōu)化，但是到了GPU和CPU上，就要差得多了。

一方面，谷歌在2018年至2021年組織和戰(zhàn)略的混亂，導(dǎo)致在對GPU進(jìn)行支持上的研發(fā)的資金不足，以及對相關(guān)問題的處理優(yōu)先級靠后。

與此同時，大概是過于專注于讓自家的TPU能在AI加速上分得更多的蛋糕，和英偉達(dá)的合作自然十分匱乏，更不用說完善對GPU的支持這種細(xì)節(jié)問題了。

另一方面，谷歌自己的內(nèi)部研究，不用想肯定都集中在TPU上，這就導(dǎo)致谷歌失去了對GPU使用的良好反饋回路。

此外，更長的調(diào)試時間、并未與Windows兼容、未跟蹤副作用的風(fēng)險等等，都增加了Jax的使用門檻以及友好程度。

現(xiàn)在，PyTorch已經(jīng)快6歲了，但完全沒有TensorFlow當(dāng)年顯現(xiàn)出的頹勢。

如此看來，想要后來者居上的話，Jax還有很長一段路要走。