谷歌丢弃 TensorFlow：豪赌 JAX

・谷歌是机器学习领域的开拓者，发布了其中一种最早的通用框架TensorFlow。

・Meta发布AI框架PyTorch之后，TensorFlow逐渐失去了开发人员的青睐。

・如今谷歌内部正押宝在一个取代TensorFlow的新AI项目JAX上。

2015年，谷歌开源了谷歌大脑（Google Brain）团队一个名为TensorFlow的小型研究项目，那年它实际上创立了现代机器学习生态系统。

TensorFlow迅速流行起来，谷歌也成为了主流AI产品的管家。

谷歌的TensorFlow曾经是无处不在的机器学习工具，此后败给了Meta的机器学习工具PyTorch。

PyTorch最初是在Facebook开发而成的，于2017年以beta的形式开放了源代码，现在日益被视为机器学习领域的领导者。

接受采访的开发人员、硬件专家、云提供商以及熟悉谷歌机器学习项目的人士给出了一致的意见。TensorFlow已经失去了开发人员的芳心。其中一些人更是放言：“PyTorch吃掉了TensorFlow 的午餐”，可谓一语中的。

专家们表示，由于自身犯下一系列战术失误、错误的开发决策以及开源社区败给Meta，谷歌在互联网上引导机器学习未来的机会可能正在逐渐消失。PyTorch已俨然成为业余开发人员和科学研究人员的首选机器学习开发工具。

如今在PyTorch的阴影下，谷歌一直在悄然构建一种名为JAX的机器学习框架，许多人将它视为是TensorFlow的接班者。JAX曾是代表“Just After eXecution”（就在执行后）的缩写词，但谷歌官方称不再代表任何意思。

接近该项目的人士告诉《商业内幕》杂志，谷歌大脑团队和谷歌的AI子公DeepMind已全面丢弃了TensorFlow，转而力推JAX，为谷歌旗下的其他公司或部门亦步亦趋铺平了道路。

谷歌代表向《商业内幕》证实，如今JAX在谷歌大脑和DeepMind几乎得到了普遍采用。

熟悉谷歌机器学习项目的人士表示，最初JAX面临来自谷歌内部的巨大阻力。

知情人士称，谷歌员工习惯于使用TensorFlow。尽管TensorFlow很笨拙，但它在谷歌员工心目当中仍不失为是统一因素，尽管令人不安。他们表示，JAX采用的方法简单得多，但改变了谷歌内部构建软件的方式。

了解该项目的人士表示，这款工具现在有望成为未来几年谷歌所有使用机器学习的产品的基础，与TensorFlow在2010年代后期的地位如出一辙。

而JAX似乎已突破了谷歌这个孤立的圈子：Salesforce告诉《商业内幕》，它已经在其研究团队中采用了JAX。

专家们经常将Julia编程语言与JAX拿来比较，Julia的创建者Viral Shah说：“JAX是了不起的技术。我认为JAX是一种单独的编程语言，恰好可以通过Python创建实例。如果你恪守JAX想要你遵守的规则，那么它就能发挥奇效，它能出色地完成它旨在完成的任务。”

谷歌现在希望再次成为领头羊，同时也从TensorFlow开发时所犯的错误中汲取教训。不过专家们表示，这是个巨大的挑战，因为现在谷歌必须取代一款已赢得了开发人员芳心的开源工具。

TensorFlow没落与PyTorch崛起

据《商业内幕》获得的数据显示，PyTorch在开发者论坛上的人气正在迅速赶超TensorFlow。Stack Overflow的人气数据显示，按论坛上所提问题的数量来衡量，TensorFlow的人气近年来停滞不前，而PyTorch的人气持续攀升。

TensorFlow起步强劲有力，推出后迅速普及开来。优步和爱彼迎等公司以及美国宇航局（NASA）等组织很快就采用了TensorFlow，开始将它用于一些最复杂的项目，这些项目需要拿海量数据集来训练算法。到2020年11月，TensorFlow已被下载了1.6亿次。

但是开发人员和熟悉该项目的人士表示，谷歌的功能蔓延（feature-creeping）和不断更新越来越使TensorFlow变得笨拙，且对用户不友好，甚至对谷歌内部的用户都不友好。随着机器学习领域迅猛发展，谷歌不得不经常使用新工具更新其框架。熟悉该项目的人士称，随着越来越多的人参与进来，这个项目在内部蔓延开来，导致人们对最初使TensorFlow成为首选工具的部分缺乏关注。

专家们告诉《商业内幕》，这种疯狂的猫捉老鼠游戏是许多率先进入市场的公司经常遇到的问题。比如说，谷歌不是第一家建立搜索引擎的公司，却能够从AltaVista或雅虎等先行者所犯的错误中汲取教训。

与此同时，PyTorch于2018年从Facebook的AI研究实验室推出了完整的生产级版本。虽然TensorFlow和PyTorch都基于机器学习专家青睐的语言Python而建，但Meta不遗余力地迎合开源社区，更是不惜投入巨资。接近TensorFlow项目的人士称，PyTorch还得益于收缩阵线、聚焦目光，而TensorFlow团队在这方面很欠缺。

机器学习初创公司Hugging Face的研究工程师Patrick von Platen说：“我们主要使用PyTorch，它得到的社区支持力度最大。我们认为PyTorch在开源方面可能做得最好。他们确保问题在线上得到回复，示例都实用有效，PyTorch还始终奉行开源优先的做法。”

一些最大的组织（包括之前依赖TensorFlow的那些组织）启动了在PyTorch上运行的项目。没过多久，特斯拉和优步等公司就在PyTorch上运行难度最大的机器学习研究项目。

添加的每项功能（有时是为了复制PyTorch备受欢迎的功能特性）使得TensorFlow对于最初的研究人员和用户受众来说越来越臃肿不堪。2017年添加“即刻执行”（eager execution）就是这样一个例子，这项原生的Python使开发人员大大简化分析和调试代码的工作。

随着PyTorch和TensorFlow一争高下，谷歌内部的一个小型研究团队致力于开发一种新框架，该框架将使用户更容易访问定制芯片（名为张量处理单元或TPU），而这种芯片是谷歌AI方法的基础，只能通过TensorFlow才可以访问。

团队研究人员包括Roy Frostig、Matthew James Johnson和Chris Leary。这三名研究人员在2018年发表了一篇题为《通过高级跟踪编译机器学习程序》的论文，论文描述了最终成为JAX的项目。

Adam Paszke之前在供职Facebook期间是PyTorch的原始开发者之一，他在2019年作为学生开始与Johnson合作，并于2020年初加入了JAX团队，成为了一名全职成员。

新项目JAX提供了一种更简单直接的设计来处理机器学习中最复杂的问题之一：将一个大问题的工作量分散到多块芯片上。JAX不是为不同的芯片运行单独的代码段，而是自动分配工作量。这种需求来自在谷歌工作的一项特别待遇：可以立即访问数量众多的TPU，完成你需要完成的任何工作。

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

谷歌内部的开发人员和研究人员借助JAX这股东风，开始采用特别研发项目。

熟悉JAX的人士表示，JAX提供了一种方法，可以避免TensorFlow对开发人员的诸多不友好之处，并迅速将复杂的技术问题分散到多个TPU上。

与此同时，PyTorch和TensorFlow都以同样的方式起家。它们起初都是研究项目，随后备受好奇，接着成为机器学习研究领域的标准，最后研究人员将它们带出学术界，进入到更广阔的天地。

然而，JAX面临几大挑战。

首先，它在很多方面仍然依赖其他框架。开发人员和专家表示，JAX没有提供轻松加载数据和预处理数据的方法，需要TensorFlow或PyTorch来处理这方面的大部分前期工作。

JAX的底层框架XLA也针对谷歌的TPU进行了高度优化。这种框架可与更传统的GPU和CPU兼容，不过接近JAX的人士表示，该项目仍有办法让GPU和CPU优化达到与TPU同等的水平。

谷歌发言人表示，重视TPU源自2018年至2021年在组织上和战略上的混乱和迷茫，结果导致投入不足，对GPU支持并未给予应有的重视度，以及缺乏与大型GPU供应商英伟达的合作，不过这几方面都在迅速改善。公司发言人表示，谷歌自己的内部研究也主要致力于TPU，这导致GPU使用方面缺乏良好的反馈回路。

Cerebras Systems是一家估值达40亿美元的初创公司，专门生产面向机器学习的大尺寸芯片。首席执行官Andrew Feldman表示，随着众多公司希望将工作量分散到许多不同类型的面向机器学习的硬件上，这种改进将至关重要。

他说：“在硬件上任何厚此薄彼的做法会立即被视为糟糕的行为，会在开源社区受到拒绝和冷落。没有人希望自己被某一家硬件供应商绑住手脚，这就是为什么机器学习框架不断涌现出来。机器学习从业人员希望确保自己的模型是可移植的，确保可以将模型带到所选择的任何硬件平台上，而不是被某一种硬件平台牢牢锁定。”

与此同时，PyTorch本身现在即将迎来第六个年头，远过了TensorFlow首次开始出现放缓迹象的时期。目前尚不清楚Meta的项目会不会遭遇与谷歌支持的前身项目（即TensorFlow）相似的命运，但这可能意味着出现新技术的时机已经成熟。几位专家和熟悉该项目的人士提到了谷歌的规模，提醒批评人士永远不要轻视这家搜索巨头。