萧箫 丰色 发自 凹非寺
量子位 | 公众号 QbitAI
收获接近16.6万个Star、见证深度学习崛起的TensorFlow,地位已岌岌可危。
并且这次,冲击不是来自老对手PyTorch,而是自家新秀JAX。
最新一波AI圈热议中,连fast.ai创始人Jeremy Howard都下场表示:
JAX正逐渐取代TensorFlow这件事,早已广为人知了。现在它就在发生(至少在谷歌内部是这样)。
LeCun更是认为,深度学习框架之间的激烈竞争,已经进入了一个新的阶段。
LeCun表示,当初谷歌的TensorFlow确实比Torch更火。然而Meta的PyTorch出现之后,现在其受欢迎程度已经超过TensorFlow了。
现在,包括Google Brain、DeepMind以及不少外部项目,都已经开始用上JAX。
典型例子就是最近爆火的DALL·E Mini,为了充分利用TPU,作者采用了JAX进行编程。有人用过后感叹:
这可比PyTorch快多了。
据《商业内幕》透露,预计在未来几年内,JAX将覆盖谷歌所有采用机器学习技术的产品。
这样看来,如今大力在内部推广JAX,更像是谷歌在框架上发起的一场“自救”。
JAX从何而来?
关于JAX,谷歌其实是有备而来。
早在2018年的时候,它就由谷歌大脑的一个三人小团队给搭出来了。
研究成果发表在了题为Compiling machine learning programs via high-level tracing的论文中:
Jax是一个用于高性能数值计算的Python库,而深度学习只是其中的功能之一。
自诞生以来,它受欢迎的程度就一直在上升。
最大的特点就是快。
一个例子感受一下。
比如求矩阵的前三次幂的和,用NumPy实现,计算需要约478毫秒。
用JAX就只需要5.54 毫秒,比NumPy快86倍。
为什么这么快?原因有很多,包括:
1、NumPy加速器。NumPy的重要性不用多说,用Python搞科学计算和机器学习,没人离得开它,但它原生一直不支持GPU等硬件加速。
JAX的计算函数API则全部基于NumPy,可以让模型很轻松在GPU和TPU上运行。这一点就拿捏住了很多人。
2、XLA。XLA(Accelerated Linear Algebra)就是加速线性代数,一个优化编译器。JAX建立在XLA之上,大幅提高了JAX计算速度的上限。
3、JIT。研究人员可使用XLA将自己的函数转换为实时编译(JIT)版本,相当于通过向计算函数添加一个简单的函数修饰符,就可以将计算速度提高几个数量级。
除此之外,JAX与Autograd完全兼容,支持自动差分,通过grad、hessian、jacfwd和jacrev等函数转换,支持反向模式和正向模式微分,并且两者可以任意顺序组成。
当然,JAX也是有一些缺点在身上的。
比如:
1、虽然JAX以加速器著称,但它并没有针对CPU计算中的每个操作进行充分优化。
2、JAX还太新,没有形成像TensorFlow那样完整的基础生态。因此它还没有被谷歌以成型产品的形式推出。
3、debug需要的时间和成本不确定,“副作用”也不完全明确。
4、不支持Windows系统,只能在上面的虚拟环境中运行。
5、没有数据加载器,得借用TensorFlow或PyTorch的。
……
尽管如此,简单、灵活又好用的JAX还是率先在DeepMind中流行起来。2020年诞生的一些深度学习库Haiku和RLax等都是基于它开发。
这一年,PyTorch原作者之一Adam Paszke,也全职加入了JAX团队。
目前,JAX的开源项目在GitHub上已有18.4k标星。
值得注意的是,在此期间,有不少声音都表示它很可能取代TensorFlow。
一方面是因为JAX的实力,另一方面主要还是跟TensorFlow自身的很多原因有关。
为什么谷歌要转投JAX?
诞生于2015年的TensorFlow,曾经也风靡一时,推出后很快超过了Torch、Theano和Caffe等一众“弄潮儿”,成为最受欢迎的机器学习框架。
然而在2017年,焕然一新的PyTorch“卷土重来”。
这是Meta基于Torch搭建的机器学习库,由于上手简单、通俗易懂,很快受到一众研究者的青睐,甚至有超过TensorFlow的趋势。
相比之下,TensorFlow却在频繁更新和界面迭代中变得越来越臃肿,逐渐失去了开发者的信任。
(从Stack Overflow上的提问占比来看,PyTorch逐年上升,TensorFlow却一直停滞不前)
在竞争之中,TensorFlow的缺点逐渐暴露出来,API不稳定、实现复杂、学习成本高等问题并没有随着更新解决多少,反而结构变得更复杂了。
相比之下,TensorFlow却没有继续发挥比较能打的“运行效率”等优势。
在学术界,PyTorch的使用率正逐渐超过TensorFlow。
尤其是在各大顶会如ACL、ICLR中,使用PyTorch实现的算法框架近几年已经占据了超过80%,相比之下TensorFlow的使用率还在不断下降。
也正是因此,谷歌坐不住了,试图用JAX夺回对机器学习框架的“主导权”。
虽然JAX名义上不是“专为深度学习构建的通用框架”,然而从发布之初起,谷歌的资源就一直在向JAX倾斜。
一方面,谷歌大脑和DeepMind逐渐将更多的库构建在JAX上。
包括谷歌大脑的Trax、Flax、Jax-md,以及DeepMind的神经网络库Haiku和强化学习库RLax等,都是基于JAX构建的。
据谷歌官方表示:
JAX生态系统开发中,也会考虑确保其与现有TensorFlow库(如Sonnet和TRFL)的设计(尽可能)保持一致。
另一方面,更多的项目也开始基于JAX实现,最近爆火的DALL·E mini项目就是其中一个。
由于能更好地利用谷歌TPU的优势,JAX在运行性能上比PyTorch要好得多,更多之前搭建在TensorFlow上的工业界项目也正在转投JAX。
甚至有网友调侃JAX如今爆火的原因:可能是TensorFlow的使用者实在无法忍受这个框架了。
那么,JAX到底有没有希望替代TensorFlow,成为与PyTorch抗衡的新力量呢?
更看好哪一个框架?
总体来看,很多人还是很坚定地站PyTorch。
他们似乎不喜欢谷歌每年都出一个新框架的速度。
“JAX虽然很吸引人,但还不够具备“革命性”的能力促使大家抛弃PyTorch来使用它。”
但看好JAX的也并非少数。
就有人表示,PyTorch是很完美,但JAX也在缩小差距。
甚至还有人疯狂给JAX打call,表示它比PyTorch要厉害10倍,并称:如果Meta不继续加把劲儿的话谷歌就会赢了。(手动狗头)
不过,总有不怎么care谁输谁赢的人,他们的目光放得很长远:
没有最好,只有更好。最重要的是更多玩家和好的idea统统都加入进来,让开源和真正优秀的创新画上等号。
项目地址:
https://github.com/google/jax
参考链接:
https://twitter.com/jeremyphoward/status/1538380788324257793
https://twitter.com/ylecun/status/1538419932475555840
https://mp.weixin.qq.com/s/AoygUZK886RClDBnp1v3jw
https://www.deepmind.com/blog/using-jax-to-accelerate-our-research
https://github.com/tensorflow/tensorflow/issues/53549
— 完 —
往期精彩回顾
适合初学者入门人工智能的路线及资料下载(图文+视频)机器学习入门系列下载中国大学慕课《机器学习》(黄海广主讲)机器学习及深度学习笔记等资料打印《统计学习方法》的代码复现专辑机器学习交流qq群955171419,加入微信群请扫码