机器学习框架局势突变：TensorFlow 逐渐式微，PyTorch 横扫顶会

自 2012 年深度学习占据人工智能领域主导地位之后，机器学习框架便成为人工智能学习必不可少的一环。经过几年时间的发展，无论是学术界还是工业界，机器学习框架都在发挥着重要的作用。近期，在 PyTorch 开发者大会上，Facebook 发布了最新版 PyTorch 1.3，并推出一系列的工具和库，就此 2019 年机器学习框架之争，仿佛瞬间进入了白热化。关于本年度两者的王者之战，你是支持学术界翘楚 PyTorch，还是站位工业界名流 TensorFlow？

本文最初发表在 The Gradient 网站，经原作者 Horace He 以及 The Gradient 网站授权，由 InfoQ 中文站翻译并分享。

自从深度学习在 2012 年重新占据主导地位以来，许多机器学习框架争相成为研究人员和从业者的新宠。从 Caffe 和 Theano 的早期学术成果，到工业界支持的大型 PyTorch 和 TensorFlow，如此多的选择，如洪水般涌来，使人们目不暇接，很难了解实际上哪个是最流行的机器学习框架。

如果你只逛 Reddit 论坛的话，你可能会有这样的错觉：每个人都在转向 PyTorch。相反，若从 Francois Chollet 的 Twitter 来判断的话，他认为 TensorFlow/Keras 可能会成为主导框架，而 PyTorch 的势头却停滞不前。

在 2019 年，机器学习框架的战争剩下两个主要竞争者：PyTorch 和 TensorFlow。我的分析表明，研究人员正在逐渐放弃 TensorFlow，陆陆续续转向 PyTorch。与此同时，在工业界中，TensorFlow 目前是首选平台，但这种情况可能并不会持续太久。

PyTorch 在学术界中日益占据主导地位

让我们来看一下数据。下图显示了 PyTorch 论文和使用 TensorFlow 或 PyTorch 的论文之间的比率。所有的线条都是向上倾斜的，并且，2019 年的每个主要会议的大多数论文是用 PyTorch 实现的。

（1） 会议图例

• CVPR、ICCV、ECCV：计算机视觉会议

• NAACL、ACL、EMNLP：自然语言处理会议

（2） 有关数据收集过程的详细情况

• 该图是通过抓取过去几年在主要机器学习会议上发表的每一篇论文而生成的。论文的分类是基于它们是否提到 PyTorch 或 TensorFlow，但不包括与 Google 或 Facebook 有关联的作者的论文，以及那些同时提到 TensorFlow 和 PyTorch 的论文。这些因素的消融可以在《机器学习框架现状：附录》（State of Frameworks: Appendix）中找到。

这些数字的交互版本可以访问：https://chillee.github.io/pytorch-vs-tensorflow/。

如果你需要更多的证据来证明 PyTorch 在学术界的发展有多快，这里有一张 PyTorch vs TensorFlow 的原始计数的图：

（注：表中，PT：PyTorch；TF：TensorFlow）

在 2018 年，PyTorch 还只是少数派。但现在，它已经成了压倒性的多数。在 CVPR，使用 PyTorch 的论文有 69%；在 NAACL 和 ACL，有 75% 以上；在 ICLR 和 ICML，有 50% 以上。虽然 PyTorch 在视觉和语言会议上的优势最为明显（分别以 2:1 和 3:1 的比率超过了 TensorFlow），但在 ICLR 和 ICML 等一般机器学习会议上，PyTorch 也比 TensorFlow 更受欢迎。

虽然有些人认为，PyTorch 仍然是一个新兴框架，试图在 TensorFlow 主导的世界中开拓一席之地，但数据却告诉我们一个不同的故事。除了 ICML 之外，甚至没有任何一个会议能让 TensorFlow 的增长跟上整体论文的增长。在 NAACL、ICLR 和 ACL 上，今年 TensorFlow 的论文实际上比去年还要少。

需要担心未来的不是 PyTorch，而是 TensorFlow。

为什么研究人员喜欢 PyTorch？

• 简单。 它与 numpy 类似，非常 Python 化，并且可以很容易地与 Python 生态系统的其他部分集成。例如，你可以简单地在 PyTorch 模型中的任何地方插入一个 pdb 断点，这样就可以起作用了。如果是 TensorFlow 的话，调试模型需要一个活动的会话，最终会使事情变得更加棘手。

• 卓越的 API。 大多数研究人员更喜欢 PyTorch 的 API，而不是 TensorFlow 的 API。部分原因是 PyTorch 的设计更好，还有部分原因是因为 TensorFlow 需要通过多次切换 API ，从而阻碍了自身的发展（如，“layers”→“slim”→“estimators”→“tf.keras”）。

• 性能。 尽管 PyTorch 的动态图提供的优化机会严格来说很少，但还是有很多传闻称，PyTorch 的速度甚至比 TensorFlow 还要快。目前尚不清楚这些传闻究竟是否真的，但至少，在这一领域中，TensorFlow 并没有取得决定性的优势。

在学术界中，TensorFlow 前景如何？

即使 TensorFlow 在功能方面上与 PyTorch 达到了同等水平，但 PyTorch 已经涵盖了社区的大多数人。这意味着，有关 PyTorch 的实现更容易被找到 ，作者们将更有动力在 PyTorch 中发布代码（这样人们就会去使用它），并且你的合作者也可能会更加喜欢 PyTorch 。因此，任何回到 TensorFlow 2.0 的迁移都可能很缓慢，如果它真发生的话。

TensorFlow 在 Google/DeepMind 中有一批忠实的粉丝，但我并不知道 Google 最终是否会放开这一点。即使是现在，Google 想要招募的许多研究人员已经在不同程度上倾向于 PyTorch，我也听到有人抱怨称，Google 内部的许多研究人员希望使用 TensorFlow 之外的框架。

此外，PyTorch 的统治地位可能会开始切断 Google 研究人员与其他研究社区的联系。他们在外部研究的基础上进行构建会比较困难，不仅如此，而且外部研究人员也不太可能在 Google 发布的代码基础上进行构建。

TensorFlow 2.0 是否能让 TensorFlow 重新吸引一些研究用户，还有待观察。尽管 Eager 模式肯定会很有吸引力，但 Keras API 就不是这样了。

用于生产环境的 PyTorch 与 TensorFlow

尽管 PyTorch 目前在学术界中占据主导地位，但快速浏览一下工业界就会发现，TensorFlow 仍然是占主导地位的框架。例如，根据 2018 年~2019 年的数据，TensorFlow 有 1541 个招聘岗位，而 PyTorch 只有 1437 个；在媒体发表的文章中，TensorFlow 有 3230 篇，而 PyTorch 有 1200 篇；在 GitHub 上，TensorFlow 有 137000 个标星，而 PyTorch 只有 7200 个标星。

那么，既然 PyTorch 已经如此受研究人员的欢迎，那为什么它在工业界上没有取得同样的成功呢？显而易见的第一个答案是：惯性。TensorFlow 比 PyTorch 早几年问世，工业界采用新技术的速度要比研究人员慢。另一个原因是 TensorFlow 在生产方面比 PyTorch 要好。但是，这意味着什么呢？

要回答这个问题，我们需要知道研究人员和工业界的需求有何不同。

研究人员关心的是他们在研究中迭代的速度能有多快，这种研究通常是在相对较小的数据集（可以在一台机器上运行的数据集）上，并且是在少于 8 个 GPU 的机器上运行的。这通常不是出于对性能方面的考虑来决定的，而是取决于他们快速实现新想法的能力。而另一方面，工业界则认为，性能是最优先考虑的。虽然运行速度提高 10% 对研究人员而言毫无意义，但这可以直接为公司节省数百万美元。另一个不同之处就是部署。研究人员将在他们自己的机器上或某个专门用于研究工作的服务器集群上进行实验。另一方面，工业界还有一连串的限制或要求。

• 没有 Python。一些公司将运行 Python 运行时开销太大而无法承受的服务器。

• 移动设备。你无法在移动设备的二进制文件中嵌入 Python 解释器。

• 服务。功能全面，如模型的无停机更新、模型之间的无缝切换、预测时间进行批处理等等。

TensorFlow 是专门围绕这些需求构建的，并针对所有这些问题提供了解决方案。图格式和执行引擎本身并不需要 Python，并且 TensorFlow Lite 和 TensorFlow Serving 分别考虑的是移动设备和服务方面。

从历史上看，PyTorch 未能满足这些考虑，因此，大多数目前在生产环境中使用 TensorFlow。

框架的“趋同化”

临近 2018 年底，发生了两件大事，让这场战争发生了翻天覆地的巨变：

1. PyTorch 引入了 JIT 编译器和 TorchScript，从而引入了基于图的特性。

2. TensorFlow 宣布，它们将在 2.0 中默认切换到 Eager 模式。

显然，这些举措都是为了解决各自的痛点。那么这些特性到底是什么，它们又能提供什么呢？

PyTorch TorchScript

PyTorch JIT 是 PyTorch 的中间表示（ Intermediate Representation，IR），称为 TorchScript。TorchScript 是 PyTorch 的图表示。你可以通过使用跟踪或脚本模式将常规的 PyTorch 模型转换为 TorchScript。跟踪采用一个函数和一个输入，记录使用该输入执行的操作，并构造 IR。虽然这一做法简单明了，但跟踪也有它的缺点。例如，它不能捕获未执行的控制流。例如，如果它执行了逻辑真块，那么它就不能捕获条件的逻辑假块。

Script 模式采用函数 / 类，重新解释 Python 代码，并直接输出 TorchScript IR。这允许它支持任意代码，但实际上，它需要重新解释 Python。

一旦 PyTorch 模型进入这个 IR 中，我们就可以获得图模式的所有好处。我们可以在 C++ 中部署 PyTorch 模型，而不需要依赖 Python，或者对其进行优化。

Tensorflow Eager

在 API 级别上，TensorFlow 的 Eager 模式本质上与 PyTorch 的 Eager 模式相同，该模式最初是由 Chainer 推出的。这为 TensorFlow 提供了 PyTorch 的 Eager 模式的大部分优点（易用性、可调试性等等）。

然而，这也给 TensorFlow 带来了同样的缺点。TensorFlow 的 Eager 模型并不能导出到非 Python 环境中，无法进行优化，也不能在移动设备上运行，等等。

这将 TensorFlow 置于与 PyTorch 相同的位置，并且它们的解析方式基本相同：你可以跟踪代码（tf.function）或重新解释 Python 代码（Autograph）。

因此，TensorFlow 的 Eager 模式并不能真正给你两全其美的体验。尽管你可以使用 tf.function 注释将 Eager 的代码转换为静态图，但这绝不是一个无缝的过程（PyTorch 的 TorchScript 也有类似的问题）。跟踪从根本上是有限的，并且重新解释 Python 代码实际上需要重写 Python 编译器的大部分内容。当然，通过限制深度学习中使用的 Python 子集，可以大大简化范围。

在默认启用 Eager 模式时，TensorFlow 会强制用户做出选择：使用 Eager 执行以简化使用，并需要重写代码以进行部署，或者完全不是用 Eager 执行。虽然这与 PyTorch 所处的情况相同，但 PyTorch 的 TorchScript 的可选属性可能比 TensorFlow 的默认的 Eager 模式更容易接受。

机器学习框架的现状

因此，我们摸清了当今机器学习框架的现状。PyTorch 占据了学术界，并试图将这一成功推广到工业界。而 TensorFlow 正试图在不牺牲太多生产能力的情况下，遏制其在学术界的流失。PyTorch 要在工业界产生有意义的影响，还有很长的一段路要走：因为 TensorFlow 过于根深蒂固，而且工业界的发展也过于缓慢。然而，从 TensorFlow 1.0 过渡到 2.0 的过程将会很困难，这就为公司评估 PyTorch 提供了一个自然的切入点。

至于它们谁将赢得未来，将取决于谁能最好地回答以下问题：

• 研究人员的偏好对工业界的影响有多大？

随着眼下这批博士研究生开始毕业时，他们会带上 PyTorch 走向岗位。这种偏好是否足够强烈，以至于公司会出于招聘目的而选择 PyTorch 呢？这些毕业生会在 PyTorch 的基础上进行创业吗？

• TensorFlow Eager 模式在可用性方面能赶上 PyTorch 吗？

问题跟踪器和在线社区给我的印象是，TensorFlow Eager 严重受到性能 / 内存问题的困扰，而且 Autograph 也有自己的问题。Google 将会花费大量的工程精力，但 TensorFlow 却被历史包袱所拖累。

• PyTorch 能以多快的速度进入生产状态？

PyTorch 仍然有许多基本问题尚未解决：没有良好的量化故事、没有移动设备、也没有服务等等。在这些问题得到解决之前，PyTorch 甚至不会成为许多公司的选择。PyTorch 能否提供一个足够有说服力的故事，让公司做出转变吗？

注：在本文发表之日，PyTorch 宣布了对量化和移动设备提供支持。这两者都还处于试验阶段，但对于 PyTorch 来说，这代表着它在这方面已经取得了重大进展。

• Google 在工业界中的孤立会对它产生伤害吗？

Google 推动 TensorFlow 的主要原因之一，是帮助其蓬勃发展的云服务。由于 Google 企图占有整个机器学习的垂直市场，这促使了与 Google 竞争的公司（Microsoft、Amazon、Nvidia）支持唯一的替代机器学习框架。

下一步是什么？

机器学习框架在多大程度上影响了机器学习研究，这一点也许没有得到人们充分的认识。它们不仅使机器学习研究成为可能，而且还使研究人员能够轻松探索的想法成为可能，并对这些想法加以限制。仅仅因为没有简单的方法在框架中表达，有多少新生的想法因此而被扼杀？PyTorch 可能已经达到了本地研究的最低要求，但是值得研究其他框架提供的内容，以及它们可能带来哪些研究机会。

高阶微分法

PyTorch 和 TensorFlow 的核心是自动微分（auto-differentiation）框架。也就是说，它们允许人们对某个函数进行求导。然而，实现自动微分的方法有很多，大多数现代机器学习框架所选择的特定实现，称为“反向模式自动微分”（reverse-mode auto-differentiation），通常被称为“反向传播”（Backpropagation）。结果证明，这种实现对于获取神经网络的导数是非常有效的。

然而，在计算高阶导数（Hessian/Hessian 向量积）时，情况就会发生变化。要有效地计算这些，需要所谓的“正向模式自动微分”（forward-mode auto-differentiation）。如果没有这种能力的话，Hessian 向量积的计算可能会慢几个数量级。

输入 Jax。Jax 是由构建原始 Autograd 的同一个人构建的，具有正向和反向模式的自动微分功能。

这使得计算高阶导数的速度比 PyTorch/TensorFlow 所能提供的速度要快几个数量级。

Jax 提供的并不仅仅是高阶导数。Jax 开发人员将 Jax 视为构成任意函数转换的框架，包括 vmap（用于自动批处理）或 pmap（用于自动并行化）。

最初的 Autograd 拥有忠实的追随者（尽管没有 GPU 支持，但 ICML 仍有 11 篇论文使用了它）。而且，很可能 Jax 很快就会建立起一个类似的专门社区，将其用于各种 n 阶导数。

代码生成

当运行 PyTorch/TensorFlow 模型时，大多数工作实际上并不是在框架本身中完成的，而是由第三方内核完成的。这些内核通常由硬件供应商提供，并由高级框架可以利用的运算符库组成。比如 MKLDNN（用于 CPU）或 cuDNN（用于 Nvidia GPU）等等。高级框架将它们的计算图分解成块，然后可以调用这些计算库。这些库代表了数千小时的工作量，并且经常针对架构和应用程序进行优化，以产生最佳性能。

最近对非标准硬件、稀疏 / 量化张量和新运算符的兴趣暴露了依赖这些操作符库的一个重大缺陷：它们不够灵活。如果你想在研究中使用新的操作者，比如胶囊网络，该怎么办呢？如果你想在一个新的硬件加速器上运行模型，但机器学习框架对此并没有提供很好的支持，那又该怎么办呢？现有的解决方案往往达不到要求。正如这篇论文最近指出的那样，胶囊网络在 GPU 上的现有实现比最佳实现要慢上两个数量级。

每一种新的硬件架构、张量类别或云算法，都大大增加了这一问题的难度。已经有许多工具可以处理不同的方面（Halide、TVM、PlaidML、Tensor Comprehensions、XLA、Taco 等等），但正确的方法仍然不清楚。

如果不投入更多的工作来解决这一问题，我们就会冒着将机器学习研究过度应用于现有工具的风险。

机器学习框架的未来

对于 TensorFlow 和 PyTorch 的未来而言，这一刻是激动人心的时刻。它们的设计已经趋同到这样的一个地步：以至于任何一个框架都不会凭借其设计而获得决定性的胜利。双方都有各自的地盘：一方占据了学术界，而另一方则占据了工业界。

就我个人而言，在 PyTorch 和 TensorFlow 之间，我赌 PyTorch 会赢得未来。机器学习仍然是一个由研究驱动的领域。工业界不能忽视研究成果，只要 PyTorch 在学术界占据主导地位，这将会迫使公司进行选择。

然而，快速发展的并不仅仅是框架。机器学习的研究本身也处于不断变化的状态。不仅框架发生了变化——五年前使用的模型 / 硬件 / 范例在今天看来，可能与我们现在所使用的截然不同。随着另一种计算模式占据主导地位，也许 PyTorch 和 TensorFlow 之间的战争将会变得无关紧要。

在所有这些相互利益冲突之中，以及围绕机器学习投入的所有资金，后退一步想想还是不错的。我们大多数人从事开发机器学习软件并不是为了赚钱，也不是为了协助公司的战略计划。我们从事机器学习的原因是，我们关心、并推进机器学习的研究，推动人工智能的民主化，或者只是仅仅关心构建一些炫酷的东西。无论你喜欢的是 TensorFlow 还是 PyTorch，我们都在努力让机器学习软件做到最好。

作者介绍：

Horace He 是康奈尔大学的学生，对编译器和机器学习的交叉领域有研究兴趣。他今年夏天在 PyTorch JIT 团队实习，但本文以及数据都是独立于 PyTorch 团队而编写 / 收集的。

The Gradient 是一本旨在普及人工智能和机器学习研究民主化的数字杂志。

原文链接：

https://thegradient.pub/state-of-ml-frameworks-2019-pytorch-dominates-research-tensorflow-dominates-industry/