深度学习框架之争：TensorFlow退守工业界，PyTorch主导学术界？

深度学习框架

Author：Horace He

Translator：MrBear

Editor：louwill

对于机器学习科研工作者和工业界从业人员来说，熟练掌握一种机器学习框架是必备技能之一。随着深度学习技术发展的突飞猛进，机器学习框架市场也渐渐度过了初期野蛮生长的阶段。大浪淘沙，目前仍然活跃的机器学习框架主要是 PyTorch 和 TensorFlow。本文从学术界和工业界两个方面深度盘点了 2019 年机器学习框架的发展趋势。

自 2012 年深度学习再度成为万众瞩目的技术以来，各种机器学习框架争相成为研究人员和从业者的新宠，可谓「你方唱罢我登场」。从早期在学术界广为使用的 Caffe 和 Theano，到业界推崇的 PyTorch 和 TensorFlow，各种各样可供选择的学习框架使得人们很难确定「谁才是真正最主流的框架」。

如果你仅仅通过浏览 Reddit 来做出判断，你可能会认为每个人都在转而使用 PyTorch；而如果你根据 Francois Chollet 推特中的内容做出判断，你会发现 TensorFlow 或 Keras 可能是主流的框架，而 PyTorch 的势头正在衰减。

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

一、PyTorch 在研究领域日益占据主导地位

让我们用数据说话！下图显示了在近些年的研究顶会中，仅仅使用了 PyTorch 框架进行研究的论文数和使用了 TensorFlow 或 PyTorch 的论文总数的比例。如图所示，每条曲线（代表不同的会议）都向上倾斜（意味着 PyTorch 的占比越来越高），而且在 2019 年的每个主要的会议中，大多数的论文都采用 PyTorch 实现。

• 会议的图例
• 数据收集过程的细节

这些图表的交互式版本链接如下：https://chillee.github.io/pytorch-vs-tensorflow/

如果你需要更多的证据来说明 PyTorch 在研究社区中获得关注的速度有多快，请看下面关于 PyTorch 和 TensorFlow 使用情况的原始统计图。

在 2018 年，PyTorch 在深度学习框架中的占比还很小。而现在，PyTorch 已成占据压倒性比重的多数。据统计，69% 的 CVPR 论文、75% 以上的 NAACL 和 ACL 论文，以及 50% 以上的 ICLR 和 ICML 论文都选择使用 PyTorch。PyTorch 在视觉和语言类的会议上（分别以 2:1 和 3:1 的比例超过了 TensorFlow）被使用的频繁度最为明显，而且 PyTorch 在 ICLR 和 ICML 等通用机器学习会议上也比 TensorFlow 更受欢迎。

虽然有些人认为 PyTorch 仍然是一个处于萌芽期的框架，试图在 Tensorflow 主导的世界中开辟出一片市场，但真实的数据却说明事实并非如此。除了在 ICML 上，其它学术会议中使用 TensorFlow 的论文的增长率甚至还赶不上整体论文数量的增长率。在 NAACL、ICLR 和 ACL 上，今年使用 TensorFlow 的论文数量实际上比去年还少。

需要担心未来发展的并不是 Pytorch，而是 TensorFlow。

1、为什么研究人员青睐 PyTorch？

• 简洁性。PyTorch 与 numpy 很相似，具有很强的 python 风格，并且很容易与 Python 生态系统中的其它组件实现集成。例如，你可以简单地在 PyTorch 模型中的任何地方添加「PDB」断点，然后就可以进行调试。在 TensorFlow 框架中，想要进行程序调试就需要一个运行中的会话，这使得调试难以进行。
• 易用的应用程序接口（API）。相较于 TensorFlow 的 API，大多数研究人员更喜欢 PyTorch 提供的 API。这在一定程度上是由于 PyTorch 的设计更好，另一方面是因为 TensorFlow 需要多次切换 API（例如，「layers」->「slim」->「estimators」->「tf.keras」）从而限制了自己的易用性。
• 卓越的性能。尽管 PyTorch 的动态图留给我们优化的机会很少，但是已经有很多有趣的报道说明 PyTorch 的运行速度和 TensorFlow 一样快（https://www.reddit.com/r/MachineLearning/comments/cvcbu6/d_why_is_pytorch_as_fast_as_and_sometimes_faster/），甚至更快（https://arxiv.org/abs/1608.07249）。目前尚不清楚这种说法是否属实，但至少，TensorFlow 在这个方面并没有获得绝对的优势。

2、TensorFlow 在研究领域的前景如何？

即使 TensorFlow 在功能方面与 PyTorch 的水平差不多，但是 PyTorch 已经拥有了研究社区中的大多数用户。这意味着我们更容易找到 PyTorch 版本的算法实现，而作者也会更有动力发布 PyTorch版本的代码（这样人们就会使用它），而你的合作者们很可能也更喜欢 PyTorch。因此，如果将代码移植回 TensorFlow 2.0 平台，这将会是一个很漫长的过程（如果真的这么做了）。

TensorFlow 在 Google/DeepMind 内部总会有一批固定的用户，但我不知道 Google 最终是否会放开这一点。即使是现在，很多 Google 想要招募的研究人员已经在不同程度上更加青睐 PyTorch 了。我也听到了一些抱怨，很多 Google 内部的研究人员希望使用 TensorFlow 以外的框架。

此外，PyTorch 的统治地位可能会开始切断 Google 研究人员与其它研究社区之间的联系。不仅 Google的研究人员将更加难以在他人研究的基础上构建自己的工作，而且外部的研究人员也不太可能基于 Google 发布的代码开展工作。

TensorFlow 2.0 是否能够为 TensorFlow 挽回一部分研究人员用户还有待观察。尽管它的动态图模式（TensorFlow 2.0 的动态图模式）一定很吸引人，但是 Keras 的 API 就并非如此了。

二、用于工业生产的 PyTorch 和 TensorFlow

虽然 PyTorch 现在在研究领域占据主导地位，但是我们快速分析一下工业界的情况就会发现，在工业界 TensorFlow 仍然是主流的框架。例如，2018 年到 2019 年的数据（参考链接：https://towardsdatascience.com/deep-learning-framework-power-scores-2018-23607ddf297a）显示，在公开招聘网站上，涉及 TensorFlow 的新招聘信息有 1541 个，而涉及 PyTorch 的新招聘信息则是 1437 个；知名科技媒体「Medium」上有 3230 篇关于 TensorFlow 的新文章，而关于 PyTorch 的新文章只有 1200 篇；在 GitHub 上，用 TensorFlow 编写的项目获得了 13,700 颗星，而用 PyTorch 编写的项目只获得了 7,200 颗星。

那么，既然 PyTorch 在研究人员中是如此受欢迎，为什么它在工业界还没有取得同样的成功呢？第一个显而易见的答案就是：惯性。TensorFlow 比 PyTorch 早诞生数年，而且工业界采用新技术的速度比研究人员要慢一些。另一个原因是：TensorFlow 比 PyTorch 更适用于生产环境。但这意味着什么呢？

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

研究人员关心的是他们在研究中迭代的速度有多快，这通常是在相对较小的数据集（可以在一台机器上运行的数据集）上、使用少于 8 个 GPU 进行的。最大的限制因素往往不是出于性能的考虑，而是他们快速实现新思路的能力。相反，工业界认为性能是需要最优先考虑的。虽然运行时的速度提升 10% 对于研究人员来说基本没有意义，但这可以直接为公司节约数百万美元的成本。

另一个区别是部署。研究人员将在他们自己的机器或某个专门用于运行研究工作的服务器集群上进行实验。另一方面，工业界在部署方面则有一连串的限制/要求：

• 不能使用 Python。对于一些公司运行的服务器来说，Python 运行时的计算开销太大了。
• 移动性。你不能在移动端的二进制文件中嵌入 Python 解释器。
• 服务性。需要满足各种需求，例如在不停机的状态下更新模型、在模型之间无缝切换、在推理时进行批处理，等等。

TensorFlow 是专门围绕这些需求构建的，并为所有这些问题提供了解决方案：计算图的版式和执行引擎本身并不需要 Python，并且通过 TensorFlow Lite 和 TensorFlow Serving 分别处理移动端和服务器端的问题。

在此前，Pytorch 还不能够很好地满足上述需求，因此大多数公司目前在生产环境下都选择使用 TensorFlow。

三、架构「趋同」

在临近 2018 年末的时候，业内发生的两件「爆炸性」事件打破了这种局面：

1. PyTorch 引入了即时（JIT，Just In Time）编译器和「TorchScript」，从而引入了基于图的特性。

2. TensorFlow 宣布它们将在 TensorFlow 2.0 版本中默认转而采用动态图模式。

显然，这些举措都旨在解决它们各自的弱点。那么，这些特性到底代表什么？它们能够为框架带来什么呢？

1、PyTorch TorchScript

PyTorch JIT可以将 PyTorch程序转换为一种名为「TorchScript」的中间表征（IR）。Torchscript 是PyTorch的「图」表征。你可以通过使用跟踪或脚模式将常规 PyTorch 模型转换为 TorchScript。跟踪接受一个函数和一个输入，记录用该输入执行的操作，并构造中间表征。虽然跟踪操作很直接，但是它也存在一些缺点。例如，它不能捕获未执行的控制流（例如，如果它执行了true情况下的程序块，它就不能捕获false 情况下的程序块。

Script模式接受一个函数/类作为输入，重新解释 Python 代码并直接输出 TorchScript 中间表征。这使得它能够支持任意代码，但它实际上需要重新解释 Python。

一旦你的 PyTorch 模型处于其中间表征状态，我们就获得了图模式的所有好处。我们可以在不依赖 Python 的情况下，在 C++ 环境中部署 PyTorch 模型，或者对其进行优化。

2、Tensorflow 动态图

在 API 的层次上，TensorFlow 的动态图模式基本上与最初由 Chainer 推崇的 PyTorch 的动态图模式相同。这为 TensorFlow 带来了PyTorch 动态图模式的大部分优势（易用性、可调试性，等等)。

然而，这也给 TensorFlow 带来了相同的缺点：TensorFlow 动态图模型不能被导出到非 python 环境中，也不能进行优化，不能在移动设备上运行，等等。

这使得 TensorFlow与 PyTorch 旗鼓相当，它们的解决方式本质上是相同的——你可以跟踪代码（tf.function）或重新解释 Python 代码（Autograph，将 print()函数和其它 Python 代码转化为纯 TensorFlow计算图代码）。

图注：TensorFlow 通过 autograph 和追踪（tracing）生成计算图

因此，TensorFlow 的动态图模式并不能真正做到「两全其美」。尽管你可以用「tf.function」注释将动态图代码转换为静态图，但这永远不会是一个无缝转换的过程（PyTorch 的 TorchScript 也有类似问题）。跟踪根本上是有限的，而重新解释 Python 代码实际上需要重写大量的 Python 编译器。当然，通过限制深度学习中用到的 Python 子集可以极大地简化这一范围。

在默认情况下启用动态图模式时，TensorFlow 使用户不得不做出选择：

• （1）为了易用性使用动态图执行，而为了进行部署需要重写函数；
• （2）完全不使用动态图执行。

尽管 PyTorch 也面临相同的问题，但相较于 TensorFlow 的「默认采用动态图」的做法， PyTorch 的 TorchScript 所具备的「选择性加入」的特性似乎让用户更愿意接受。

四、机器学习框架的现状

以上原因造就了机器学习框架领域当今的局面。PyTorch 拥有研究人员的市场，并试图将这种成功延伸至工业界。而 TensorFlow 则试图在不牺牲过多的生产能力的情况下，阻止其在研究社区中所占市场份额的流失。PyTorch 想要在工业界产生巨大的影响肯定还有很长的路要走，毕竟 TensorFlow 在工业界已经根深蒂固，而且工业界革新的速度也相对较慢。然而，从 TensorFlow 1.0 过渡到 2.0 将是一个艰难的过程，这也让公司自然而然地会评估是否采用 PyTorch。

在未来，哪种框架能够「笑到最后」取决于以下问题：

• 研究人员的倾向会对工业界产生多大的影响？随着当下的博士生纷纷毕业，他们将会把使用 PyTorch 的习惯延续新的岗位上去。这种倾向是否足够令公司选择招聘使用 PyTorch 的雇员呢？毕业生们会基于 Pytorch 技术创业吗？
• TensorFlow 的动态图模式能否在易用性方面追赶上 PyTorch？我从对该问题紧密追踪的人以及在线社区哪里感受到的情况是，TensorFlow 的动态图严重受到「性能/内存」问题的困扰，而且「Autograph」自身也存在许多问题。Google 将为此付出大量的工程努力，但是 TensorFlow 还是会受到许多「历史遗留问题」的困扰。
• PyTorch 能多快在生产环境中被大规模采用？PyTorch 还有许多基本问题有待解决，比如没有好的量化方式、不能满足移动性和服务性需求。对于大多数公司来说，在这些问题被妥善解决之前，它们甚至都不会考虑使用 PyTorch。PyTorch 是否能足够令公司信服，改变使用的机器学习框架呢？（注：近日，PyTorch 宣布了支持量化和移动性功能，这两种功能尚处于试验阶段，但代表了 PyTorch 在这方面取得了重大进展。）
• Google 在业内被孤立会让 TensorFlow 受挫吗？Google 推动 TensorFlow 的主要原因之一是帮助其迅速发展的云服务。由于 Google 正在试图占领整个机器学习垂直市场，这刺激其它竞争公司（微软、亚马逊、英伟达）纷纷转向唯一的机器学习框架备选方案——PyTorch。

五、机器学习框架之争的下半场将如何上演？

我们还没有充分认识到机器学习框架对机器学习的研究产生了多大的影响。它们不仅使机器学习研究可以进行下去，还将一些研究的思路进行了限制，使这些思路切实可行，让研究人员能够很容易地对这些思路进行探索。事实上，有多少新奇的想法仅仅因为不能在某种框架中用一种简单的方式表达出来而破灭？目前看来，PyTorch 可能已经能够达到研究的「局部最小值」，但是我们仍然需要看看其它的框架提供了哪些特性，还存在哪些研究的机会。

1、高阶微分

PyTorch 和 Tensorflow 的核心是自动微分框架。也就是说，这些框架使我们可以对某些函数进行求导。然而，尽管有许多方法可以实现自动微分，但大多数现代机器学习框架采用的都是「逆向模式自动微分」（通常又被称为「反向传播」）算法。事实证明，这种方法对于对神经网络进行求导是极为高效的。

然而，在计算高阶导数（Hessian/Hessian 向量积）时，情况就不一样了。想要高效地计算这些值需要用到「前向模式自动微分」。不使用这个功能的话，对 Hessian 向量积的计算速度会慢几个量级。

接下来我们将介绍「Jax」。Jax 是由开发原始「Autograd」的人员开发的，它同时具备「前向传播」和「逆向传播」自动微分。这使得我们对于高阶导数的计算相较于 PyTorch/TensorFlow 可以提供的方法快了几个数量级。

然而，Jax 并不仅仅只提供了计算高阶导数的方法。Jax 的开发者将其视为一种组成任意函数变换的框架，包括「vmap」（针对自动化批处理）或「pmap」（针对自动化并行计算）。

原始的 Autograd 拥有其忠实的拥趸（尽管没有 GPU 的支持，仍然有 11 篇 ICML 上发表的论文采用了它），Jax 可能很快就会拥有一个类似的忠实用户社区，使用它来求解任何类型的 n 阶导数。

2、代码生成

当你运行一个 PyTorch/TensorFlow 模型时，大部分工作实际上并不是在框架本身中完成的，而是由第三方内核完成的。这些内核通常由硬件供应商提供，由高级框架可以利用的算子库组成，例如，MKLDNN（用于 CPU）或 cuDNN（用于英伟达 GPU）。高级框架将计算图分解成块，它可以调用上面提到的计算库。构建这些库需要数千个人时的工作量，并针对架构和应用程序进行了优化以获得最佳性能。

然而，最近研究人员对于「非标准硬件」、「稀疏/量子化张量」和「新运算符」的关注暴露了依赖这些运算符库的一个缺陷：它们并不灵活。

如果你想在研究中使用像胶囊网络这样的新算子，你该怎么做？如果你想在机器学习框架目前还不支持的新型硬件加速器上运行你的模型，你又该怎么做？现有的解决方案往往还不够完善。正如论文「Machine Learning Systems are Stuck in a Rut」（论文地址：https://dl.acm.org/citation.cfm?id=3321441）所提到的，现有的胶囊网络在 GPU 上的实现比最优的实现慢了两个数量级。

每种新的硬件架构、张量或算子的类别，都大大提高了该问题的难度。目前已经有许多处理工具（如 Halide、TVM、PlaidML、TensorComprehensions、XLA、Taco 等）可以处理各种问题，但是真正正确的处理方法还有待探索。

如果不能进一步解决这个问题，我们就会有一定风险将我们的机器学习研究与我们拥有的研究工具「过度拟合」。

六、机器学习框架的未来

这对于 TensorFlow 和 PyTorch 的未来而言是激动人心的时刻：它们的设计逐渐趋同，它们都不太可能凭借其设计获得决定性的胜利。与此同时，这两种机器学习框架都有其各自主导的领域——PyTorch 在学术界占据主导，而 TensorFlow 在工业界则更受欢迎。

就我个人而言，在 PyTorch 和 TensorFlow 之间，我认为 PyTorch 更有胜算。机器学习仍然是一个由研究驱动的领域。工业界不能忽视科学研究的成果，只要 PyTorch 在研究领域占据主导地位，就会迫使公司转而使用 PyTorch.

然而，不仅机器学习框架迭代得非常快，机器学习研究本身也处于一场巨大的变革之中。发生变化的不仅仅是机器学习框架，5 年后使用的模型、硬件、范式与我们现在使用的可能有非常大的区别。也许，随着另一种计算模型占据主导地位，PyTorch 与 TensorFlow 之间的机器学习框架之争也将烟消云散。

置身于这些错综复杂的利益冲突、以及投入在机器学习领域的大量资金中，退一步，也许海阔天空。大多数从事机器学习软件的工作不是为了赚钱，也不是为了协助公司的战略计划，而是想要推进机器学习的研究，关心人工智能民主化，也或许他们只是想创造一些很酷的东西。我们大多数人并不是为了赚钱或协助我们企业的战略而从事机器学习软件事业，我们从是机器学习工作的原因只是因为——我们关心机器学习研究的发展，使人工智能走进千家万户，或者仅仅只是因为我们向创造一些很酷的东西。无论你更喜欢 TensorFlow 还是 PyTorch，我们都怀着一个共同的目的：尽力做出最棒的机器学习软件！