很喜歡有些網(wǎng)友的一句話：

「這孩子實(shí)在不行，咱再要一個(gè)吧?！?/p>

谷歌還真這么干了。

養(yǎng)了七年的TensorFlow終于還是被Meta的PyTorch干趴下了，在一定程度上。

谷歌眼見(jiàn)不對(duì)，趕緊又要了一個(gè)——「JAX」，一款全新的機(jī)器學(xué)習(xí)框架。

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

TensorFlow的黃昏和PyTorch的崛起

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

當(dāng)時(shí)，TensorFlow只是Google Brain的一個(gè)小項(xiàng)目。

誰(shuí)也沒(méi)有想到，剛一問(wèn)世，TensorFlow就變得非?；鸨?/p>

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

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

不過(guò)，谷歌好像并沒(méi)有十分在乎這么多用戶(hù)的感受。

奇奇怪怪的界面和頻繁的更新都讓TensorFlow對(duì)用戶(hù)越來(lái)越不友好，并且越來(lái)越難以操作。

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

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

于是，越來(lái)越多的人加入了這個(gè)項(xiàng)目，導(dǎo)致整個(gè)團(tuán)隊(duì)慢慢失去了重點(diǎn)。

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

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

這種困局對(duì)最先打入某一市場(chǎng)的公司來(lái)說(shuō)是避不開(kāi)的。

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

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

正是因?yàn)樯厦孢@些原因，原先給谷歌賣(mài)命的開(kāi)發(fā)者，慢慢對(duì)老東家失去了信心。

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

2017年，PyTorch的測(cè)試版開(kāi)源。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

根據(jù)所應(yīng)用的情況，JAX會(huì)自動(dòng)地將若干個(gè)芯片組合而成一個(gè)小團(tuán)體，而不是讓一個(gè)去單打獨(dú)斗。

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

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

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

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

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

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

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

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

舉個(gè)栗子：

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

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

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

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

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

此外，XLA和Autograd可以任意組合，甚至可以利用pmap方法一次使用多個(gè)GPU或TPU內(nèi)核進(jìn)行編程。

而將JAX與Autograd和Numpy相結(jié)合的話，就可以獲得一個(gè)面向CPU、GPU和TPU的易于編程且高性能的機(jī)器學(xué)習(xí)系統(tǒng)了。

顯然，谷歌這一次吸取了教訓(xùn)，除了在自家全面鋪開(kāi)以外，在推進(jìn)開(kāi)源生態(tài)的建設(shè)方面，也是格外地積極。

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

縱觀整場(chǎng)「明爭(zhēng)暗斗」，賈揚(yáng)清表示，在批評(píng)TensorFlow的進(jìn)程中，AI系統(tǒng)認(rèn)為Pythonic的科研就是全部需求。

但一方面純Python無(wú)法實(shí)現(xiàn)高效的軟硬協(xié)同設(shè)計(jì)，另一方面上層分布式系統(tǒng)依然需要高效的抽象。

而JAX正是在尋找更好的平衡，谷歌這種愿意顛覆自己的pragmatism非常值得學(xué)習(xí)。

causact R軟件包和相關(guān)貝葉斯分析教科書(shū)的作者表示，自己很高興看到谷歌從TF過(guò)渡到JAX，一個(gè)更干凈的解決方案。

谷歌的挑戰(zhàn)

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

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

除了各種隱藏的bug以外，JAX在一些問(wèn)題上仍然要依賴(lài)于其他框架。

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

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

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

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

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

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

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

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

如此看來(lái)，想要后來(lái)者居上的話，Jax還有很長(zhǎng)一段路要走。