收獲接近16.6萬個(gè)Star、見證深度學(xué)習(xí)崛起的TensorFlow,地位已岌岌可危。
并且這次,沖擊不是來自老對手PyTorch,而是自家新秀JAX。
最新一波AI圈熱議中,連fast.ai創(chuàng)始人Jeremy Howard都下場表示:
JAX正逐漸取代TensorFlow這件事,早已廣為人知了。現(xiàn)在它就在發(fā)生(至少在谷歌內(nèi)部是這樣)。
LeCun更是認(rèn)為,深度學(xué)習(xí)框架之間的激烈競爭,已經(jīng)進(jìn)入了一個(gè)新的階段。
LeCun表示,當(dāng)初谷歌的TensorFlow確實(shí)比Torch更火。然而Meta的PyTorch出現(xiàn)之后,現(xiàn)在其受歡迎程度已經(jīng)超過TensorFlow了。
現(xiàn)在,包括Google Brain、DeepMind以及不少外部項(xiàng)目,都已經(jīng)開始用上JAX。
典型例子就是最近爆火的DALL·E Mini,為了充分利用TPU,作者采用了JAX進(jìn)行編程。有人用過后感嘆:
這可比PyTorch快多了。
據(jù)《商業(yè)內(nèi)幕》透露,預(yù)計(jì)在未來幾年內(nèi),JAX將覆蓋谷歌所有采用機(jī)器學(xué)習(xí)技術(shù)的產(chǎn)品。
這樣看來,如今大力在內(nèi)部推廣JAX,更像是谷歌在框架上發(fā)起的一場“自救”。
JAX從何而來?
關(guān)于JAX,谷歌其實(shí)是有備而來。
早在2018年的時(shí)候,它就由谷歌大腦的一個(gè)三人小團(tuán)隊(duì)給搭出來了。
研究成果發(fā)表在了題為Compiling machine learning programs via high-level tracing的論文中:
Jax是一個(gè)用于高性能數(shù)值計(jì)算的Python庫,而深度學(xué)習(xí)只是其中的功能之一。
自誕生以來,它受歡迎的程度就一直在上升。
最大的特點(diǎn)就是快。
一個(gè)例子感受一下。
比如求矩陣的前三次冪的和,用NumPy實(shí)現(xiàn),計(jì)算需要約478毫秒。
用JAX就只需要5.54 毫秒,比NumPy快86倍。
為什么這么快?原因有很多,包括:
1、NumPy加速器。NumPy的重要性不用多說,用Python搞科學(xué)計(jì)算和機(jī)器學(xué)習(xí),沒人離得開它,但它原生一直不支持GPU等硬件加速。
JAX的計(jì)算函數(shù)API則全部基于NumPy,可以讓模型很輕松在GPU和TPU上運(yùn)行。這一點(diǎn)就拿捏住了很多人。
2、XLA。XLA(Accelerated Linear Algebra)就是加速線性代數(shù),一個(gè)優(yōu)化編譯器。JAX建立在XLA之上,大幅提高了JAX計(jì)算速度的上限。
3、JIT。研究人員可使用XLA將自己的函數(shù)轉(zhuǎn)換為實(shí)時(shí)編譯(JIT)版本,相當(dāng)于通過向計(jì)算函數(shù)添加一個(gè)簡單的函數(shù)修飾符,就可以將計(jì)算速度提高幾個(gè)數(shù)量級。
除此之外,JAX與Autograd完全兼容,支持自動差分,通過grad、hessian、jacfwd和jacrev等函數(shù)轉(zhuǎn)換,支持反向模式和正向模式微分,并且兩者可以任意順序組成。
當(dāng)然,JAX也是有一些缺點(diǎn)在身上的。
比如:
1、雖然JAX以加速器著稱,但它并沒有針對CPU計(jì)算中的每個(gè)操作進(jìn)行充分優(yōu)化。
2、JAX還太新,沒有形成像TensorFlow那樣完整的基礎(chǔ)生態(tài)。因此它還沒有被谷歌以成型產(chǎn)品的形式推出。
3、debug需要的時(shí)間和成本不確定,“副作用”也不完全明確。
4、不支持Windows系統(tǒng),只能在上面的虛擬環(huán)境中運(yùn)行。
5、沒有數(shù)據(jù)加載器,得借用TensorFlow或PyTorch的。
……
盡管如此,簡單、靈活又好用的JAX還是率先在DeepMind中流行起來。2020年誕生的一些深度學(xué)習(xí)庫Haiku和RLax等都是基于它開發(fā)。
這一年,PyTorch原作者之一Adam Paszke,也全職加入了JAX團(tuán)隊(duì)。
目前,JAX的開源項(xiàng)目在GitHub上已有18.4k標(biāo)星,比TensorFlow高了不少了。
值得注意的是,在此期間,有不少聲音都表示它很可能取代TensorFlow。
一方面是因?yàn)镴AX的實(shí)力,另一方面主要還是跟TensorFlow自身的很多原因有關(guān)。
為什么谷歌要轉(zhuǎn)投JAX?
誕生于2015年的TensorFlow,曾經(jīng)也風(fēng)靡一時(shí),推出后很快超過了Torch、Theano和Caffe等一眾“弄潮兒”,成為最受歡迎的機(jī)器學(xué)習(xí)框架。
然而在2017年,煥然一新的PyTorch“卷土重來”。
這是Meta基于Torch搭建的機(jī)器學(xué)習(xí)庫,由于上手簡單、通俗易懂,很快受到一眾研究者的青睞,甚至有超過TensorFlow的趨勢。
相比之下,TensorFlow卻在頻繁更新和界面迭代中變得越來越臃腫,逐漸失去了開發(fā)者的信任。
(從Stack Overflow上的提問占比來看,PyTorch逐年上升,TensorFlow卻一直停滯不前)
在競爭之中,TensorFlow的缺點(diǎn)逐漸暴露出來,API不穩(wěn)定、實(shí)現(xiàn)復(fù)雜、學(xué)習(xí)成本高等問題并沒有隨著更新解決多少,反而結(jié)構(gòu)變得更復(fù)雜了。
相比之下,TensorFlow卻沒有繼續(xù)發(fā)揮比較能打的“運(yùn)行效率”等優(yōu)勢。
在學(xué)術(shù)界,PyTorch的使用率正逐漸超過TensorFlow。
尤其是在各大頂會如ACL、ICLR中,使用PyTorch實(shí)現(xiàn)的算法框架近幾年已經(jīng)占據(jù)了超過80%,相比之下TensorFlow的使用率還在不斷下降。
也正是因此,谷歌坐不住了,試圖用JAX奪回對機(jī)器學(xué)習(xí)框架的“主導(dǎo)權(quán)”。
雖然JAX名義上不是“專為深度學(xué)習(xí)構(gòu)建的通用框架”,然而從發(fā)布之初起,谷歌的資源就一直在向JAX傾斜。
一方面,谷歌大腦和DeepMind逐漸將更多的庫構(gòu)建在JAX上。
包括谷歌大腦的Trax、Flax、Jax-md,以及DeepMind的神經(jīng)網(wǎng)絡(luò)庫Haiku和強(qiáng)化學(xué)習(xí)庫RLax等,都是基于JAX構(gòu)建的。
據(jù)谷歌官方表示:
JAX生態(tài)系統(tǒng)開發(fā)中,也會考慮確保其與現(xiàn)有TensorFlow庫(如Sonnet和TRFL)的設(shè)計(jì)(盡可能)保持一致。
另一方面,更多的項(xiàng)目也開始基于JAX實(shí)現(xiàn),最近爆火的DALL·E mini項(xiàng)目就是其中一個(gè)。
由于能更好地利用谷歌TPU的優(yōu)勢,JAX在運(yùn)行性能上比PyTorch要好得多,更多之前搭建在TensorFlow上的工業(yè)界項(xiàng)目也正在轉(zhuǎn)投JAX。
甚至有網(wǎng)友調(diào)侃JAX如今爆火的原因:可能是TensorFlow的使用者實(shí)在無法忍受這個(gè)框架了。
那么,JAX到底有沒有希望替代TensorFlow,成為與PyTorch抗衡的新力量呢?
更看好哪一個(gè)框架?
總體來看,很多人還是很堅(jiān)定地站PyTorch。
他們似乎不喜歡谷歌每年都出一個(gè)新框架的速度。
“JAX雖然很吸引人,但還不夠具備“革命性”的能力促使大家拋棄PyTorch來使用它!
但看好JAX的也并非少數(shù)。
就有人表示,PyTorch是很完美,但JAX也在縮小差距。
甚至還有人瘋狂給JAX打call,表示它比PyTorch要厲害10倍,并稱:如果Meta不繼續(xù)加把勁兒的話谷歌就會贏了。(手動狗頭)
不過,總有不怎么care誰輸誰贏的人,他們的目光放得很長遠(yuǎn):
沒有最好,只有更好。最重要的是更多玩家和好的idea統(tǒng)統(tǒng)都加入進(jìn)來,讓開源和真正優(yōu)秀的創(chuàng)新畫上等號。