JAX 正从谷歌的边缘项目走向核心

谷歌于 2015 年开放了一种小型资源，即 2015 年谷歌大脑研究项目，名为 TensorFlow，实质上创造了现代机器学习的生态系统。其普及速度之快，令谷歌成为了主流人工智能产品的拥有者。

但是现在，情况已经完全不同了，谷歌已经丧失了开发者的人心，转向了 Meta。

谷歌的 TensorFlow 曾经是一款无所不在的机器学习工具，但后来却落后于 Meta 机器学习工具 PyTorch。PyTorch 首次在 Facebook 开发，在 2017 年作为测试版开源，PyTorch 逐渐成为了领军人物。

在采访开发者、硬件专家、云提供商以及熟悉谷歌机器学习工作的人士时，他们的观点也是相同的。TensorFlow 在争夺开发者人心的竞争中落败。其中有些人甚至使用了令人难以理解的确切说法：“PyTorch 正在享用 TensorFlow 的午餐”。

专家称，经过 Meta 在开源社区中的一系列战术失误、开发决策和智胜策略，谷歌引领未来互联网上的机器学习的机遇可能将会慢慢消逝。自那以后，PyTorch 就成了开发人员和科学研究人员的机器学习开发工具。

谷歌在 PyTorch 的阴影下，悄悄地研发了一种 名为 JAX（一度是“Just After eXecution”的缩写，但是正式意义上已经没有什么意义了）的机器学习框架，许多人认为它是 TensorFlow 的继承者。

该项目的知情人士向 Insider 透露，谷歌大脑和谷歌的 DeepMind 人工智能公司基本上都放弃了 TensorFlow，转而改用 JAX，以便为谷歌其他部门的跟进铺平道路。一位谷歌代表向 Insider 证实，JAX 目前已经在谷歌大脑和 DeepMind 中得到了广泛的采用。

项目地址：https://github.com/google/jax/releases

熟悉谷歌机器学习工作的人士称，JAX 起初受到了很大的内部阻力。他们表示，谷歌员工过去一直使用 TensorFlow。尽管它使用起来可能很困难，但是在谷歌的雇员中，这始终是一个让人困扰的统一因素。JAX 的方法更加简单得多，但是他们说，这也改变了谷歌在内部构建软件的方式。

熟悉该项目的人表示，这款工具将会有望在今后数年内成为所有使用机器学习的谷歌产品的支柱，就像 2010 年代末 TensorFlow 所做的那样。

而 JAX 似乎已经从谷歌的边缘走向了中心。Salesforce 告诉 Insider，她已经在自己的研究团队中采用了 JAX。

“JAX 是一项工程壮举，”Julia 编程语言的创建者 Viral Shah 说，专家们经常将其与 JAX 相提并论。“我认为 JAX 是一种独立的编程语言，通过 Python 进行实例化。如果你坚持 JAX 想要的规则，它就能发挥它的神奇力量，而且可以做到让人叹为观止。”

谷歌现在想要重回霸主的位置，同时也从开发 TensorFlow 时犯下的错误中汲取教训。但是，专家们说，由于他现在不得不放弃这个已经赢得了开发者人心的开源工具，这将是一项很大的挑战。

TensorFlow 的衰落与 PyTorch 的崛起

据提供给 Insider 的数据显示，PyTorch 在一家必读的开发者论坛上的帖子正在迅速追赶 TensorFlow。来自 Stack Overflow 的参与度数据表明，以其在论坛问题中的份额衡量，TensorFlow 在最近几年中的人气一直处于停滞状态，而 PyTorch 的参与度却在持续上升。

TensorFlow 一开始就有很强劲的势头，并且随着其发布而越来越流行。像 Uber 和 Airbnb 这样的公司，以及像 NASA 这样的组织很快就选择了 TensorFlow，并将 TensorFlow 应用到更复杂的项目中，这些项目要求在巨大的数据集上训练算法。截至 2020 年 11 月，TensorFlow 已累计下载 1.6 亿次。

但是，谷歌不断增量的功能更新使 TensorFlow 变得很不方便，而且对用户来说也很不友好，即使是谷歌内部的雇员、开发者和接近该项目的人也会这么认为。随着机器学习领域以惊人的速度发展，谷歌必须 经常使用新的工具来更新其框架。熟悉该项目的人士表示，由于更多的人参与进来，这个项目已经开始向公司内部扩展，而不是将重点放在那些原本让 TensorFlow 成为首选工具的部分。

专家告诉 Insider，这种狂热的猫鼠游戏在很多率先推出的公司中是很常见的。举例来说，谷歌并非首家创立搜索引擎的公司，它只是能够 从 AltaVista 或雅虎等祖先的错误中学习。

与此同时，PyTorch 在 2018 年在 Facebook 人工智能研究实验室 发布了其完整版本。尽管 TensorFlow 和 PyTorch 都是基于 Python，而 Python 是机器学习专家的首选语言，但是 Meta 已经在迎合开源社区的需求方面投入了大量资金。据了解 TensorFlow 项目的人说，PyTorch 在一定程度上得益于专注小而美的产品，而 TensorFlow 团队一直在错过。

“我们主要使用 PyTorch；它拥有最广泛的社区支持，”机器学习创业公司 Hugging Face 的研究工程师 Patrick von Platten 说。“我们认为 PyTorch 可能在开源方面做得最好。他们确保问题能够得到在线解答。所有的例子都有效。PyTorch 一直有一个非常开源的方法。”

有些最大的组织（包括那些依赖 TensorFlow 的组织），Weave 项目是在 PyTorch 上运行 的 。不久之后，像特斯拉和 Uber 这样的公司 开始在 PyTorch 上进行最艰难的机器学习研究项目。

每一个额外的特性，有时会复制那些让 PyTorch 非常受欢迎的元素，使得 TensorFlow 对其最初的研究人员和用户而言，变得越来越臃肿了。其中一个例子是，它在  2017 年增加了一个“敏锐执行”（Keen Execution）的特性，这是 Python 的原生特性，可以让开发者很容易对自己的代码进行分析和调试。

进入 JAX，谷歌机器学习的未来

随着 PyTorch 和 TensorFlow 之间的战斗爆发，谷歌内部的一支小型研究团队致力于开发一种新的框架，以便更容易访问专门设计的芯片（称为张量处理单元，或 TPU），这些芯片是其人工智能方法的基础，只能通过 TensorFlow 获得。

团队研究人员包括 Roy Frostige、Matthew James Johnson 和 Leary 在 2018 年发布了一篇题为《通过高级可追溯性编译机器学习软件》（Compilation of machine learning software through high-level traceability）的论文，描述了最终成为 JAX 的内容。

Adam Paszky 是之前在 Facebook 工作期间 PyTorch 的原作者之一，他在 2019 年开始以学生身份与 Johnson 合作，并在 2020 年初全职加入 JAX 团队。

新项目 JAX 提出了一种更加直观的设计，他可以处理最复杂的机器学习问题：将一个大问题的工作分散到多个芯片上。JAX 不是为不同的芯片运行单一的代码位，而是自动分配工作。这个要求来自在谷歌工作的一个很大的特点：只要你需要，就可以立即使用大量的 TPU 来完成任何你想要的一切。

JAX 解决了谷歌研究人员在处理需要越来越多计算能力的大型问题时所面临的一个基本问题。

JAX Wind Catch，这个 skunkworks 项目正在被谷歌内部的开发者和研究人员所接受。熟悉该项目的人士说，这是一种绕开许多开发者对 TensorFlow 不敬的做法，并且能够很快地把复杂的技术问题分散到多个 TPU 上。

谷歌在 JAX 方面的最大挑战是用 PyTorch 实施其 Meta 战略

同时，PyTorch 和 TensorFlow 都是以同样的方式开始的。它们一开始是研究项目，然后是好奇心，接着成为机器学习研究中的标准。最后，研究人员将它们从学术界传播到了世界其他地方。

但是，JAX 也面临着许多挑战。首先，在许多方面，它对其他框架的依赖性还很强。开发者和专家们表示，JAX 没有提供一种加载数据和预处理数据的方法，它需要 TensorFlow 或 PyTorch 来处理大量的设置。

JAX 基本框架 XLA，它也为谷歌 TPU 设备进行了极大的优化。该框架还能与更传统的 GPU 和 CPU 协同工作，不过，据了解该项目的人表示，该项目仍有办法通过优化 GPU 和 CPU，从而达到与 TPU 的同等水平。

谷歌发言人表示，对 TPU 的关注源于 2018 年至 2021 年的监管和战略上的混乱，造成了缺乏投资和对 GPU 支持的次优优先级，以及缺乏与大型 GPU 提供商 Nvidia 的合作，这两方面都在迅速改善。这位发言人说，谷歌自己的内部研究也非常关注 TPU，导致缺乏良好的 GPU 使用反馈。

Cerebras Systems 公司首席执行官 Andrew Feldman 说，随着各公司希望将其工作分散到不同类型的机器上，专注于机器学习，因此这种改进将成为未来的关键。这家市值 40 亿美元的初创公司，制造专注于机器学习的大型芯片。

他说：“任何以一种设备为特色而非另一种设备的行为，都会立刻被认为是一种恶劣的行为，并且会受到开源社区的排斥。没有人愿意被限制在一个单一的硬件提供商，这就是机器学习框架出现的原因。机器学习从业者希望确保他们的模型是可移植的，他们可以将其移植到他们选择的任何硬件平台上，而不是被锁定在一个平台上。”

同时，PyTorch 本身现在已经有将近 6 年的历史了，远远超过了 TensorFlow 首次开始出现放缓迹象的年龄。目前还不清楚 Meta 项目是否会面临与其谷歌支持的前辈类似的命运，但这可能意味着新事物出现的时机已经成熟。有些专家和接近该项目的人士指出了谷歌的规模过大，并告诫批评家千万别把希望寄托在这个搜索巨头身上。

原文链接：

https://comnavitottori.com/how-pytorch-beat-tensorflow-forcing-google-to-bet-on-jax/