放弃TensorFlow，谷歌全面转向JAX

OpenCV学堂 2022-06-15 16:00

TI MCU方案：能源基础设施实时控制 【TI资料】基于新型C29内核的MCU技术资料

点击上方↑↑↑“OpenCV学堂”关注我

来源：公众号机器之心授权

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