谷歌在框架上发起的一场“自救”
作者：萧箫 丰色
来源：量子位公众号 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
往期推荐
腾讯云开发者
