TensorFlow，危！抛弃者正是谷歌自己

　　欢迎关注“新浪科技”的微信订阅号：techsina 

　　文/萧箫

　　来源：量子位（ID: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统统都加入进来，让开源和真正优秀的创新画上等号。