TensorFlow,危!抛弃者正是谷歌自己

萧箫 丰色 发自 凹非寺
量子位 | 公众号 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高了不少了。

值得注意的是,在此期间,有不少声音都表示它很可能取代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

版权所有,未经授权不得以任何形式转载及使用,违者必究。