放弃TensorFlow，谷歌全面转向JAX

作者丨Matthew Lynley

来源丨机器之心

TensorFlow 大概已经成为了谷歌的一枚「弃子」。

2015 年，谷歌大脑开放了一个名为「TensorFlow」的研究项目，这款产品迅速流行起来，成为人工智能业界的主流深度学习框架，塑造了现代机器学习的生态系统。

七年后的今天，故事的走向已经完全不同：谷歌的 TensorFlow 失去了开发者的拥护，因为他们已经转向了 Meta 推出的另一款框架 PyTorch。

曾经无处不在的机器学习工具 TensorFlow 已经悄悄落后，而 Facebook 在 2017 年开源的 PyTorch 正在成为这个领域的霸主。

近日，外媒 Business Insider 采访了一系列开发人员、硬件专家、云供应商以及与谷歌机器学习工作关系密切的人，获得了同样的观点。TensorFlow 已经输掉了这场战争，其中有人用了一个鲜明的比喻：「PyTorch 吃掉了 TensorFlow 的午餐。」

专家们表示，鉴于战术失误、开发决策和 Meta 在开源社区中的一系列智取策略，谷歌引领互联网机器学习未来的机会正在逐渐消失。

在 PyTorch 的阴影下，谷歌正在悄悄地开发一个机器学习框架，就是 JAX（曾是「Just After eXecution」的首字母缩写，但官方说法中不再代表任何东西），许多人将其视为 TensorFlow 的继承者。

接近该项目的人士告诉 Insider，谷歌大脑和 DeepMind 在很大程度上放弃了 TensorFlow，转而使用 JAX。这为谷歌的其他部门铺平了跟随的道路，一位谷歌内部人士向 Insider 证实，JAX 现在几乎已在谷歌大脑和 DeepMind 中被全球采用。

接近谷歌机器学习工作的人士表示，最初 JAX 面临着来自内部的强烈反对，一些人认为谷歌员工已经习惯了使用 TensorFlow。尽管它可能很难用，但它一直是谷歌员工中的统一因素。他们说，JAX 方法是要简单得多，但它会改变 Google 内部构建软件的方式。

熟悉该项目的人士表示，Jax 现在有望成为未来几年所有使用机器学习的谷歌产品的支柱，就像 TensorFlow 在 2015 年之后几年所做的那样。

「JAX 是一项工程壮举，」Julia 编程语言创建者 Viral Shah 说。「我认为 JAX 是一种通过 Python 实例化的独立编程语言。如果你遵守 JAX 想要的规则，它就可以发挥它的魔力，这真是令人惊叹。」

现在，谷歌希望在这场竞赛中再次获得金牌，同时也从开发 TensorFlow 时所犯的错误中吸取教训，但这将是一个巨大的挑战。

TensorFlow 的暮光，PyTorch 的崛起

根据提供给 Insider 的数据，PyTorch 在一些必读开发者论坛上的帖子正在迅速赶超 TensorFlow。Stack Overflow 的参与度数据显示，以论坛问题份额衡量， TensorFlow 受欢迎程度近年来停滞不前，而 PyTorch 的参与度继续上升。

TensorFlow 起步强劲，推出后受欢迎程度不断提高。Uber 和 Airbnb 等公司以及 NASA 等机构很快就开始将其用于一些复杂的项目，这些项目需要在大量数据集上训练算法。截至 2020 年 11 月，TensorFlow 已被下载 1.6 亿次。

但谷歌持续及增量的功能更新使得 TensorFlow 变得笨拙，且对用户不友好，即使是谷歌内部的那些人、开发人员和与项目关系密切的人都认为如此。随着机器学习领域以惊人的速度发展，谷歌不得不经常使用新工具更新其框架。接近该项目的人士表示，该项目已经在内部传播开来，越来越多的人参与其中，不再专注最初是什么让 TensorFlow 成为首选工具。

专家告诉 Insider，对于许多拥有引领者身份的公司来说，这种疯狂的猫鼠游戏是一个反复出现的问题。例如，谷歌并不是第一家建立搜索引擎的公司，它能够从 AltaVista 或 Yahoo 等前辈的错误中吸取教训。

2018 年，PyTorch 推出了完整版。虽然 TensorFlow 和 PyTorch 都建立在 Python 之上，但 Meta 在满足开源社区的需求方面投入了大量资金。熟悉 TensorFlow 项目的人士说，PyTorch 还受益于专注做一些 TensorFlow 团队错过的事情。

「我们主要使用 PyTorch，它拥有最多的社区支持，」机器学习初创公司 Hugging Face 的研究工程师 Patrick von Platten 说。「我们认为 PyTorch 可能在开源方面做得最好，他们能确保在线回复问题，所有示例都能 work。」

一些最大的组织机构开始在 PyTorch 上运行项目，包括那些曾经依赖 TensorFlow 的机构。不久之前，特斯拉、Uber 等公司就在 PyTorch 上运行了他们最艰巨的机器学习研究项目。

TensorFlow 的新增功能有时会复制使 PyTorch 流行的元素，使得 TensorFlow 对于其最初的研究人员和用户受众来说越来越臃肿。一个这样的例子是它在 2017 年增加了「敏锐执行」，这是 Python 的原生特性，使开发人员可以轻松分析和调试他们的代码。

尝试用 JAX 自救

随着 PyTorch 和 TensorFlow 之间竞争日益激烈，谷歌内部的一个小型研究团队开发了一个新框架 JAX，该框架将更容易访问张量处理单元（TPU）——一种谷歌专门为机器学习和 TensorFlow 定制的芯片。

团队研究人员 Roy Frostige、Matthew James Johnson 和 Chris Leary 在 2018 年发表了一篇名为《Compilation of machine learning software through high-level traceability》的论文，介绍了这个新框架 JAX。PyTorch 的原始作者之一 Adam Paszky 于 2020 年初全职加入 JAX 团队。

论文地址：https://cs.stanford.edu/~rfrostig/pubs/jax-mlsys2018.pdf

JAX 提出了一种直接的设计来解决机器学习中最复杂的问题之一：将一个大问题的工作分散到多个芯片上。JAX 不会为不同的芯片运行单段代码，而是自动分配工作，即时访问尽可能多的 TPU，以满足运行需要。这解决了谷歌研究人员对算力的巨大需求。

PyTorch 和 TensorFlow 都是以同样的方式开始的，首先是研究项目，然后成为机器学习研究的标准化工具，从学界扩散到更多地方。

然而，JAX 在很多方面仍然依赖于其他框架。开发人员表示：JAX 没有加载数据和预处理数据的简单方法，需要使用 TensorFlow 或 PyTorch 进行大量数据处理。

JAX 基础框架 XLA 也针对谷歌 TPU 设备做了很大的优化。该框架还适用于更传统的 GPU 和 CPU。了解该项目的人称：该项目仍有办法对 GPU 和 CPU 进行优化，以达到 TPU 同等水平。

谷歌发言人表示，2018 年至 2021 年谷歌内部对 GPU 的支持策略欠优，缺乏与大型 GPU 供应商英伟达的合作，因此谷歌转而关注 TPU，谷歌自己内部的研究也主要集中在 TPU 上，导致缺乏良好的 GPU 使用反馈循环。从这个意义上说，谷歌对 TPU 的重视和专注也属迫不得已。

然而，如 Cerebras Systems 的 CEO Andrew Feldman 所说：「任何以一种设备优于另一种设备的行为都是不良做法，会遭到开源社区的抵制。没有人希望局限于单一的硬件供应商，这就是机器学习框架应运而生的原因。机器学习从业者希望确保他们的模型是可移植的，可以移植到他们选择的任何硬件上，而不是被锁定在一个平台上。」

如今，PyTorch 快 6 岁了，TensorFlow 在这个年纪早已出现衰落的迹象。也许 PyTorch 有一天也会被新框架取代，这尚未可知。但至少，新框架出现的时机已经成熟。

原文链接：https://www.businessinsider.com/facebook-pytorch-beat-google-tensorflow-jax-meta-ai-2022-6

本文仅做学术分享，如有侵权，请联系删文。

3D视觉工坊精品课程官网：3dcver.com

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缺陷检测教程：理论、源码与实战

重磅！3DCVer-学术论文写作投稿 交流群已成立

扫码添加小助手微信，可申请加入3D视觉工坊-学术论文写作与投稿 微信交流群，旨在交流顶会、顶刊、SCI、EI等写作与投稿事宜。

同时也可申请加入我们的细分方向交流群，目前主要有3D视觉、CV&深度学习、SLAM、三维重建、点云后处理、自动驾驶、多传感器融合、CV入门、三维测量、VR/AR、3D人脸识别、医疗影像、缺陷检测、行人重识别、目标跟踪、视觉产品落地、视觉竞赛、车牌识别、硬件选型、学术交流、求职交流、ORB-SLAM系列源码交流、深度估计等微信群。

一定要备注：研究方向+学校/公司+昵称，例如：”3D视觉 + 上海交大 + 静静“。请按照格式备注，可快速被通过且邀请进群。原创投稿也请联系。

▲长按加微信群或投稿

▲长按关注公众号

3D视觉从入门到精通知识星球：针对3D视觉领域的视频课程（三维重建系列、三维点云系列、结构光系列、手眼标定、相机标定、激光/视觉SLAM、自动驾驶等）、知识点汇总、入门进阶学习路线、最新paper分享、疑问解答五个方面进行深耕，更有各类大厂的算法工程人员进行技术指导。与此同时，星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息，打造成集技术与就业为一体的铁杆粉丝聚集区，近4000星球成员为创造更好的AI世界共同进步，知识星球入口：

学习3D视觉核心技术，扫描查看介绍，3天内无条件退款

 圈里有高质量教程资料、答疑解惑、助你高效解决问题

觉得有用，麻烦给个赞和在看~