首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >谷歌在框架上发起的一场“自救”

谷歌在框架上发起的一场“自救”

作者头像
用户1737318
发布2022-06-24 18:17:59
6520
发布2022-06-24 18:17:59
举报
文章被收录于专栏:人工智能头条人工智能头条

作者:萧箫 丰色

来源:量子位公众号 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,更像是谷歌在框架上发起的一场“自救”。

PART 01

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自身的很多原因有关。

PART 02

为什么谷歌要转投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抗衡的新力量呢?

PART 03

更看好哪一个框架?

总体来看,很多人还是很坚定地站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

END

往期推荐

AI越进化越跟人类大脑像!Meta找到了机器的“前额叶皮层”,AI学者和神经科学家都惊了

大一女生废话编程爆火!懂不懂编程的看完都拴Q了

微软再曝“丑闻”:在办公室看 VR 黄片,“HoloLens 之父”即将离职!

本文参与 腾讯云自媒体分享计划,分享自微信公众号。
原始发表:2022-06-23,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 人工智能头条 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体分享计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档