作者丨拉燕如愿 好困
来源新智元
谷歌Meta之争看来还没完！TensorFlow干不过还有JAX，能否战胜二战？PyTorch？
我非常喜欢一些网民的话：
「这孩子真的不行，我们再要一个吧。」
谷歌真的这么做了。
养了七年的TensorFlow终于还是被Meta的PyTorch干趴下，在一定程度上。
谷歌看错了，赶紧又要了一个——「JAX」，全新的机器学习框架。
最近超级热DALL·E Mini大家都知道，它的模型是基于JAX编程充分利用谷歌TPU优势。
TensorFlow的黄昏和PyTorch的崛起
2015年，谷歌开发了机器学习框架——TensorFlow问世。
当时，TensorFlow只是Google Brain一个小项目。
没人想到，刚出来，TensorFlow变得很热。
优步、爱彼迎等大公司在用，NASA这种国家机构也在使用。它们也用于他们最复杂的项目。
截至2020年11月，TensorFlow下载次数已达到1.6亿次。
然而，谷歌似乎并不太在意这么多用户的感受。
奇怪的界面和频繁的更新都让TensorFlow对用户越来越不友好，操作也越来越困难。
就连谷歌内部也觉得这个框架在走下坡路。
事实上，谷歌如此频繁的更新也是无奈的，毕竟只有这样才能赶上机器学习领域的快速迭代。
因此，越来越多的人加入了这个项目，导致整个团队慢慢失去了重点。
而原本让TensorFlow那些成为首选工具的闪光点，也埋藏在无数元素中，不再被重视。
这种现象被Insider形容为一种「猫鼠游戏」。公司就像猫，迭代的新需求就像老鼠。猫要时刻保持警惕，随时扑向老鼠。
这种困境是第一个进入某个市场的公司无法避免的。
比如谷歌索引擎而言，谷歌并不是第一个。所以谷歌可以从前辈开始（AltaVista、Yahoo等等)总结经验的失败，应用于自身的发展。
可惜到了TensorFlow谷歌在这里被困住了。
正是因为这些原因，原本为谷歌努力的开发者慢慢对老东家失去了信心。
过去无处不在TensorFlow渐渐倒下，输了Meta的后起之秀——PyTorch。
2017年，PyTorch测试版开源。
2018年，Facebook发布了人工智能研究实验室PyTorch完整版本。
值得一提的是，PyTorch和TensorFlow都是基于Python开发的，而Meta更注重维护开源社区，甚至投入大量资源。
而且，Meta关注谷歌的问题，认为不能重蹈覆辙。他们专注于一小部分功能，并尽最大努力。
Meta谷歌没有后尘。这个首先在Facebook开发的框架逐渐成为行业标杆。
一家机器学习初创公司的研究工程师说，「我们基本上都用PyTorch。它的社区和开源是最好的。不仅要回答问题，还要给出实用的例子。」
面对这种情况，谷歌的开发者、硬件专家、云提供商和任何与谷歌机器学习相关的人员在接受采访时都说了同样的话。他们认为TensorFlow失去了开发者的心。
经历了一系列的明争暗斗，Meta最终占上风。
有专家表示，谷歌未来继续引领机器学习的机会正在慢慢流失。
PyTorch逐渐成为普通开发者和研究人员的首选工具。
从Stack Overflow根据开发者论坛提供的互动数据PyTorch越来越多的问题，关于TensorFlow近年来一直处于停滞状态。
就连文章开头提到的优步等公司也转向了PyTorch了。
甚至，PyTorch后来的每一次更新，似乎都在玩TensorFlow的脸。
谷歌机器学习的未来——JAX
就在TensorFlow和PyTorch打得热火朝天的时候，谷歌内部的一个「小黑马研究团队」开始开发一个新的框架，可以更方便地使用TPU。
一篇题为2018年的文章《Compiling machine learning programs via high-level tracing》的论文，让JAX项目浮出水面，作者是Roy Frostig、Matthew James Johnson和Chris Leary。
从左到右依次是这三位大神
而后，PyTorch原作者之一Adam Paszke，2020年初全职加入JAX团队。
JAX机器学习中最复杂的问题之一是多核处理器调度问题。
根据应用情况，JAX它会自动将几个芯片组合成一个小组，而不是让一个人独自战斗。
这样做的好处是让它尽可能多TPU你可以在一瞬间得到反应，从而燃烧我们「炼丹小宇宙」。
最后，与臃肿相比TensorFlow，JAX解决了谷歌内心的一个大问题:如何快速访问？TPU。
下面简单介绍一下构成JAX的Autograd和XLA。
Autograd主要用于梯度优化，可自动区分Python和Numpy代码。
它既可以用来处理Python子集，包括循环、递归和闭包，也可以求导导数的导数。
此外，Autograd支持梯度的反向传输，这意味着它可以有效地获得标量值函数相对于数组值参数的梯度，以及前向模式的微分，两者可以任意组合。
XLA（Accelerated Linear Algebra）可以加速TensorFlow模型不需要更改源代码。
当一个程序运行时，所有操作都由执行器单独执行。每个操作都有预编译GPU内核实现后，执行器将分配到内核实现。
举个栗子：
def model_fn(x, y, z): return tf.reduce_sum(x y * z)在没有XLA运行时，这部分将启动三个核心：一个用于乘法，一个用于加法，另一个用于减法。
而XLA通过加法、乘法和减法「融合」到单个GPU从而内核。
这种集成操作不会将内存产生的中间值写入y*z内存x y*z；相反，它直接将这些中间计算的结果「流式传输」同时，将它们完全保存给用户GPU中。
在实践中，XLA性能提升的7倍左右，5倍左右batch大小改进。
此外，XLA和Autograd它可以任意组合，甚至可以使用pmap多种方法一次使用GPU或TPU编程内核。
而将JAX与Autograd和Numpy如果结合起来，你可以得到一个面向CPU、GPU和TPU机器学习系统易于编程，性能高。
显然，谷歌这一次吸取了教训，除了在自家全面铺开以外，在推进开源生态的建设方面，也是格外地积极。
2020年DeepMind正式投入JAX怀抱，这也宣布谷歌亲自结束，从此各种开源库层出不穷。
纵观整场「明争暗斗」，贾扬清说，他在批评TensorFlow的进程中，AI系统认为Pythonic科研是所有需求。
但一方面纯Python另一方面，上层分布式系统仍然需要高效的抽象来实现高效的软硬协同设计。
而JAX谷歌正在寻找更好的平衡，谷歌愿意颠覆自己pragmatism值得学习。
causact R软件包和贝叶斯分析教科书的作者说，他很高兴看到谷歌从TF过渡到JAX，一个更干净的解决方案。
谷歌的挑战
作为新秀，Jax虽然可以借鉴PyTorch和TensorFlow这两位老前辈的优点，但有时后发也可能带来劣势。
首先，JAX还太「年轻」，作为一个实验框架，远未达到谷歌成熟产品的标准。
除各种隐藏bug以外，JAX在某些问题上仍然依赖于其他框架。
需要使用加载和预处理数据TensorFlow或PyTorch处理大部分设置。
显然，这与理想有关「一站式」框架相距甚远。
其次，JAX主要针对TPU高度优化，但到了GPU和CPU上，差得多。
一方面，歌在2018年至2021年组织和战略的混乱，导致在对GPU进行支持上的研发的资金不足，以及对相关问题的处理优先级靠后。
与此同时，大概是过于专注于让自家的TPU能在AI加速上分得更多的蛋糕，和英伟达的合作自然十分匮乏，更不用说完善对GPU的支持这种细节问题了。
另一方面，谷歌自己的内部研究，不用想肯定都集中在TPU上，这就导致谷歌失去了对GPU使用的良好反馈回路。
此外，更长的调试时间、并未与Windows兼容、未跟踪副作用的风险等等，都增加了Jax的使用门槛以及友好程度。
现在，PyTorch已经快6岁了，但完全没有TensorFlow当年显现出的颓势。
如此看来，想要后来者居上的话，Jax还有很长一段路要走。
参考资料：
https://www.businessinsider.com/facebook-pytorch-beat-google-tensorflow-jax-meta-ai-2022-6
本文仅做学术分享，如有侵权，请联系删文。
1.面向自动驾驶领域的多传感器数据融合技术
2.面向自动驾驶领域的3D点云目标检测全栈学习路线！(单模态+多模态/数据+代码)3.彻底搞透视觉三维重建：原理剖析、代码讲解、及优化改进4.国内首个面向工业级实战的点云处理课程5.激光-视觉-IMU-GPS融合SLAM算法梳理和代码讲解6.彻底搞懂视觉-惯性SLAM：基于VINS-Fusion正式开课啦7.彻底搞懂基于LOAM框架的3D激光SLAM: 源码剖析到算法优化8.彻底剖析室内、室外激光SLAM关键算法原理、代码和实战(cartographer+LOAM +LIO-SAM)
9.从零搭建一套结构光3D重建系统[理论+源码+实践]
10.单目深度估计方法：算法梳理与代码实现
11.自动驾驶中的深度学习模型部署实战
12.相机模型与标定(单目+双目+鱼眼）
13.重磅！四旋翼飞行器：算法与实战
14.ROS2从入门到精通：理论与实战
15.国内首个3D缺陷检测教程：理论、源码与实战
扫码添加小助手微信，可
也可申请加入我们的细分方向交流群，目前主要有、、、、、等微信群。
一定要备注：，例如：”3D视觉 + 上海交大 + 静静“。请按照格式备注，可快速被通过且邀请进群。也请联系。
▲长按加微信群或投稿
▲长按关注公众号
：针对3D视觉领域的五个方面进行深耕，更有各类大厂的算法工程人员进行技术指导。与此同时，星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息，打造成集技术与就业为一体的铁杆粉丝聚集区，近4000星球成员为创造更好的AI世界共同进步，知识星球入口：
学习3D视觉核心技术，扫描查看介绍，3天内无条件退款
圈里有高质量教程资料、答疑解惑、助你高效解决问题