TensorFlow与PyTorch模型部署性能比较

前言

2022了，选 PyTorch 还是 TensorFlow？之前有一种说法：TensorFlow 适合业界，PyTorch 适合学界。这种说法到 2022 年还成立吗？从模型可用性、部署便捷度和生态系统三个方面对比了两个框架的优缺点，针对不同身份的读者给出了不同的选择建议。

模型可用性对比：PyTorch 胜出
在模型可用性方面，PyTorch 和 TensorFlow 表现出了明显的区别。都有官方模型库，但是对于从业者来说，这些库里的模型可能还不够。因此，有必要比较一下这两个框架官方模型库之外的模型来源是否丰富。
HuggingFace
HuggingFace 的存在使得深度学习从业者仅借助几行代码就能将训练、微调好的 SOTA 模型整合到pipeline 中。下图显示了两个框架的 HuggingFace 模型可用情况。从图中看出，HuggingFace 中大约有 85% 的模型只能在 PyTorch 上用，剩下的模型还有一半，可以在 PyTorch 上用。相比之下，只有 16% 的模型能在 TensorFlow 上用，只有 8% 是 TensorFlow 所独有的。

 如果把范围缩小到 HuggingFace 上最受欢迎的 30 个模型，可以看到类似的结果。在这 30 个模型中，能在 TensorFlow 上用的还不到 2/3，但能在 PyTorch 上用的却达到了 100%，没有哪个模型只能在 TensorFlow 上用。

 研究论文

对于研究者来说，从最近发表的论文中获取模型是非常重要的，可以专注于其他重要的工作。尝试在不同的框架中重新创建新的模型会浪费宝贵的时间。
考虑到 PyTorch 实际上是一个研究型框架，预期在 HuggingFace 上观察到的趋势会蔓延至整个社区，结果不出所料。
如下图所示，汇总了 8 个顶级研究期刊上的论文在过去几年里的框架采用情况。结果显示，PyTorch 的采用率增长迅速，几年时间就从原来的 7% 长到了近 80%。

 PyTorch 的这种增长势头很大程度上是拜 TensorFlow 所赐。很多转向 PyTorch 的研究者都表示 TensorFlow 1 太难用了。尽管 2019 年发布的 TensorFlow 2 改掉了一些问题，但PyTorch 的增长势头已经难以遏制。

如果看一下那些中途换框架的研究者的比例，可以得出类似的结论。如下图所示，在 2018 年还在用 TensorFlow 的论文作者中，有 55% 的人在 2019 年转向了 PyTorch，但 2018 年就在用 PyTorch 的人有 85% 都留了下来。

 Papers with Code

最后看一下 Papers with Code 网站。在这一部分，绘制了使用 PyTorch、TensorFlow 或其他框架的论文随时间变化的比例。从中可以看出，使用 PyTorch 的论文在稳步增长：在本季度创建的 4500 个库中，有 60% 是在 PyTorch 中实现的，只有 11% 是在 TensorFlow 中实现的。相比之下，TensorFlow 的使用率在稳步下降，2019 年 TensorFlow 2 的发布也没有扭转这一趋势。

 小结

从以上数据可以明显看出，PyTorch 目前在研究领域占据主导地位。虽然 TensorFlow 2 解决了研究者使用该框架进行研究的一些痛点，但 PyTorch 却没有给研究者回头的理由。此外，TensorFlow 两大版本间的向后兼容性问题只会让这种趋势愈演愈烈。
当然，这里还有几个例外需要注意：
Google AI：谷歌发布的论文自然会用 TensorFlow。鉴于在论文方面谷歌比 Facebook 更高产，一些研究者可能会发现掌握 TensorFlow 还是很有用的。
DeepMind：DeepMind 也用 TensorFlow，比 Facebook 高产。创建了一个名叫 Sonnet 的 TensorFlow 高级 API，用于研究目的。有人称这个 API 叫「科研版 Keras」，那些考虑用 TensorFlow 做研究的人可能会用到。此外，DeepMind 的 Acme 框架可能对于强化学习研究者很有用。
OpenAI：OpenAI 在 2020 年宣布了全面拥抱 PyTorch 的决定。但之前的强化学习基线库都是在 TensorFlow 上部署的。基线提供了高质量强化学习算法的实现，因此 TensorFlow 可能还是强化学习从业者的最佳选择。
JAX：谷歌还有另一个框架——JAX，在研究社区中越来越受欢迎。与 PyTorch 和 TensorFlow 相比，JAX 的开销要小得多。但同时，JAX 和前两个框架差别也很大，因此迁移到 JAX 对于大多数人，可能不是一个好选择。目前，有越来越多的模型/论文已经在用 JAX，但未来几年的趋势依然不甚明朗。
所以总体来看，在第一轮（模型可用性）的对比中，PyTorch 完胜。
部署便捷性对比：TensorFlow 胜出
从一开始，TensorFlow 就是一个面向部署的首选框架，因为有一系列可以提高端到端深度学习效率的工具，如 TensorFlow Serving 和 TensorFlow Lite。
相比之下，PyTorch 在部署方面一度表现平平，但近年来，在努力缩小这一差距。2020年推出的 TorchServe 和2021年推出的 PyTorch Live 为用户提供了急需的本地部署工具。但二者之间的差距还有多大？看进一步的分析。
TensorFlow
TensorFlow 有静态图，可以优化推理性能。当使用 TensorFlow 部署模型时，可以根据具体应用选择使用 TensorFlow Serving 或 TensorFlow Lite。
TensorFlow Serving：
TensorFlow Serving 用于在服务器上部署 TensorFlow 模型，无论是在内部还是在云上，在 TensorFlow Extended（TFX）端到端机器学习平台中使用。Serving 使得用模型标记（model tag）将模型序列化到定义良好的目录中变得很容易，可以选择在保持服务器架构和 API 静态的情况下，使用哪个模型来进行推理请求。
Serving 可以帮用户轻松地在 gRPC 服务器上部署模型，这些服务器运行谷歌为高性能 RPC 打造的开源框架。gRPC 的设计意图是连接不同的微服务生态系统，这些服务器非常适合模型部署。Serving 通过 Vertex AI 和 Google Cloud 紧密地集成在一起，与Kubernetes 以及 Docker 进行了集成。
TensorFlow Lite：
TensorFlow Lite 用于在移动或物联网 / 嵌入式设备上部署 TensorFlow 模型。TFLite 对这些设备上的模型进行了压缩和优化，解决了设备上的 AI 的 5 个约束——延迟、连接、隐私、大小和功耗。可以使用相同的 pipeline，同时导出基于标准 Keras 的 SavedModels（和 Serving 一起使用）和 TFLite 模型，这样就能比较模型的质量。
TFLite 可用于 Android、iOS、微控制器和嵌入式 Linux。TensorFlow 针对 Python、Java、C++、JavaScript 和 Swift 的 API，为开发人员提供了广泛的语言选项。
PyTorch
PyTorch 在简化部署方面有所投入。以前，PyTorch 用户需要使用 Flask 或 Django 在模型上构建一个 REST API，但现在有了 TorchServe 和 PyTorch Live 的本地部署选项。
TorchServe：
TorchServe 是 AWS 和 Facebook 合作的开源部署框架，于 2020 年发布。具有端点规范、模型归档和指标观测等基本功能，但仍然不如 TensorFlow。TorchServe 同时支持 REST 和 gRPC API。
PyTorch Live：
PyTorch 于 2019 年首次发布 PyTorch Mobile，旨在为部署优化的机器学习模型创建端到端工作流，适用于 Android、iOS 和 Linux。
PyTorch Live 于 12 月初发布，以移动平台为基础。使用 JavaScript 和 React Native，创建带有相关 UI 的跨平台 iOS 和 Android AI 应用。设备上的推理仍然由 PyTorch Mobile 执行。Live 提供了示例项目来辅助入门，计划在未来支持音频和视频输入。
小结
目前，TensorFlow 依然在部署方面占有优势。Serving 和 TFLite 比 PyTorch 的同类型工具要稳健一些。将 TFLite 与谷歌的 Coral 设备一起用于本地 AI 的能力是许多行业的必备条件。相比之下，PyTorch Live 只专注于移动平台，TorchServe 仍处于起步阶段。因此综合来看，第二轮（部署便捷性）的胜出者是 TensorFlow。
如果既想用 TensorFlow 的部署基础设施，又想访问只能在 PyTorch 中使用的模型，推荐使用 ONNX 将模型从 PyTorch 移植到 TensorFlow。
生态系统对比：TensorFlow 胜出
2022 年，将 PyTorch 和 TensorFlow 分开考虑，一个重要的因素是所处的生态系统不同。PyTorch 和 TensorFlow 都提供了易于部署、管理、分布式训练的工具，从建模的角度讲都是能力很强的框架。相比之下，在生态系统方面的差异更重要。
PyTorch 
Hub：
PyTorch Hub 作为面向研究的官方平台，用于与预训练模型共享存储库。Hub 拥有广泛类别的模型，包括用于音频、视觉、NLP 任务的模型，还有用于生成任务的 GAN 模型。
SpeechBrain：
SpeechBrain 是 PyTorch 的官方开源语音工具包。SpeechBrain 能够完成自动语音识别（ASR）、说话人识别、验证和分类等任务。如果不想构建任何模型，想要一个具有情感分析、实体检测等功能的即插即用工具，可以选择使用 AssemblyAI 的 Speech-to-Text API。
当然，PyTorch 的工具页面还有很多其他有用的库，包括为计算机视觉和自然语言处理量身定制的库，如 fast.ai。
TorchElastic：
TorchElastic 是 AWS 和 Facebook 2020 年联合发布的分布式训练工具，可管理工作进程并协调重启行为，以便用户在计算节点集群上训练模型，这些节点可以动态变化而不会影响训练。因此，TorchElastic 可防止因服务器维护或网络问题等导致的灾难性故障，不会丢失训练进度。TorchElastic 具有与 Kubernetes 集成的特性，已集成到 PyTorch 1.9+ 中。
TorchX：
TorchX 是一个用于快速构建和部署机器学习应用程序的 SDK。TorchX 包括 Training Session Manager API，可在支持的调度程序上启动分布式 PyTorch 应用程序。TorchX 负责启动分布式作业，原生支持由 TorchElastic 局部管理的作业。
Lightning：
PyTorch Lightning 有时被称为 PyTorch 的 Keras。虽然这种类比不准确，但 Lightning 的确是简化 PyTorch 中模型工程和训练过程的有用工具，自 2019 年首次发布，已经逐渐趋于成熟。Lightning 以面向对象的方式处理建模过程，定义了可重用和可跨项目使用的可共享组件。
TensorFlow
Hub：
TensorFlow Hub 是一个经过训练的机器学习模型库，可以进行微调，只需几行代码就能使用像 BERT 这样的模型。Hub 包含适用于不同用例的 TensorFlow、TensorFlow Lite 和 TensorFlow.js 模型，可用于图像、视频、音频和文本处理。
Model Garden：
如果现成的预训练模型不适用于用户的应用，TensorFlow 的存储库 Model Garden 可以提供 SOTA 模型的源代码。对于想要深入了解模型工作原理，或根据需要修改模型的用户，Model Garden 将非常有用。
Model Garden 包含谷歌维护的官方模型、研究人员维护的研究模型和社区维护的精选社区模型。TensorFlow 的长期目标是在 Hub 上提供来自 Model Garden 的模型的预训练版本，使 Hub 上的预训练模型在 Model Garden 中具有可用的源代码。
Extended（TFX）：
TensorFlow Extended 是 TensorFlow 用于模型部署的端到端平台。该平台的功能强大，包括：加载、验证、分析和转换数据；训练和评估模型；使用 Serving 或 Lite 部署模型；跟踪 artifact 及其依赖项。TFX 还可以与 Jupyter 或 Colab 一起使用，可以使用 Apache Airflow/Beam 或 Kubernetes 进行编排。TFX 与 Google Cloud 紧密集成，可与 Vertex AI Pipelines 一起使用。
Vertex AI：
Vertex AI 是 Google Cloud 2021年刚刚发布的统一机器学习平台，旨在统一 GCP、AI Platform 和 AutoML，成为一个平台。Vertex AI 能够以无服务器方式编排工作流，帮助用户自动化、监控和管理机器学习系统。Vertex AI 可以存储工作流的 artifact，让用户可以跟踪依赖项和模型的训练数据、超参数和源代码。
Coral：
尽管有各种各样的 SaaS 公司依赖基于云的人工智能，但许多行业对本地人工智能的需求也在不断增长，Google Coral 就是为了满足这一需求创建的。Coral 是一个完整的工具包，可以使用本地 AI 构建产品。Coral 于 2020 年发布，解决了部署部分 TFLite 中提到的实现板载 AI 的问题，克服了隐私和效率等方面的困难。
Coral 提供了一系列用于原型设计、生产和传感的硬件产品，一些本质上是增强型的树莓派，专为 AI 应用程序创建，能够利用 Edge TPU 在低功耗设备上进行高性能推理。Coral 提供用于图像分割、姿态估计、语音识别等任务的预编译模型，为希望创建本地 AI 系统的开发人员提供支持。创建模型的基本步骤如下面的流程图所示。

 TensorFlow.js：

TensorFlow.js 是一个用于机器学习的 JavaScript 库，允许用户使用 Node.js 在浏览器和服务器端训练和部署模型。
Cloud：
TensorFlow Cloud 是一个可以将本地环境连接到 Google Cloud 的库，API 旨在弥补本地机器上模型构建和调试与 GCP 上分布式训练和超参数调整间的差距，而无需使用 Cloud Console。
Colab：
Google Colab 是一个基于云的 notebook 环境，与 Jupyter 非常相似。Colab 易于连接到 Google Cloud 进行 GPU 或 TPU 训练，Colab 可以和 PyTorch 一起使用。
Playground：
Playground 是一个小而精致的可视化工具，用于帮助用户理解神经网络的基础知识。要户可以更改 Playground 内置神经网络的层数和大小，实时查看神经网络是如何学习特征的，用户还可以看到改变学习率和正则化强度等超参数，如何影响不同数据集的学习过程。Playground 允许实时播放学习过程，高度直观的方式查看输入在训练过程中是如何转换的。Playground 提供了一个开源的小型神经网络库，是自身的构建基础，用户能够查看其源代码的具体细节。
Datasets：
谷歌研究院的 Datasets 是谷歌定期发布的数据集的整合资源。谷歌还提供了数据集搜索，访问更广泛的数据集资源。当然，PyTorch 用户可以利用这些数据集。
小结
总的来说，尽管 PyTorch 和 TensorFlow 有很多相似和共享的资源，但在生态系统方面，终究是 TensorFlow 更胜一筹。谷歌投入巨资确保深度学习的每个相关领域都有完善的产品。与 Google Cloud 和 TFX 的紧密集成，使端到端的开发过程变得轻而易举，将模型移植到 Google Coral 设备的便利性，让 TensorFlow 在某些行业取得了压倒性的胜利。
我应该选 PyTorch 还是 TensorFlow？
正如期望的那样，PyTorch 与 TensorFlow 没有决出明确的胜负。某一个框架在特定用例方面是优于另一个框架的。为了帮助读者做出选择，作者汇总了一些建议。在下面的流程图中，每个图表都针对不同的兴趣领域量身定制。
如果在业界怎么办？

 如果在工业界执行深度学习工程任务，很可能会使用 TensorFlow，应该坚持使用。对于需要将模型投入生产的人来说，TensorFlow 强大的部署框架和端到端的 TensorFlow Extended 平台是很珍贵的。能在 gRPC 服务器上进行轻松部署以及模型监控和工件跟踪是行业应用的关键。

鉴于 PyTorch 最近发布了 TorchServe，如果需要访问仅在 PyTorch 中可用的 SOTA 模型，可以考虑使用 PyTorch。在这种情况下，使用 ONNX 在 TensorFlow 的部署工作流中部署转换后的 PyTorch 模型。
如果正在构建移动应用，鉴于 PyTorch 最近发布了 PyTorch Live，可以考虑使用 PyTorch。如果需要音频或视频输入，应该使用 TensorFlow。如果正在构建使用 AI 的嵌入式系统或 IoT 设备，鉴于 TFLite + Coral 生态系统，仍然应该使用 TensorFlow。
总之，在工业界，如果必须选择一种框架，请选择 TensorFlow。
如果是研究者怎么办？

 如果是一名研究人员，大概率会使用 PyTorch，应该坚持使用，大多数 SOTA 模型都适用于 PyTorch。

当然，这条规则有几个值得注意的例外，最值得注意的是：强化学习领域的一些研究应该考虑使用 TensorFlow。TensorFlow 有一个用于强化学习的原生 Agents 库，DeepMind 的 Acme 框架是在 TensorFlow 中实现的。OpenAI 的 baseline 模型存储库在 TensorFlow 中实现的，尽管 OpenAI 的 Gym 在 TensorFlow 或 PyTorch 中都能使用。如果打算使用 TensorFlow 进行研究，应该查看 DeepMind 的 Sonnet 以获得更高级别的抽象。
如果不想使用 TensorFlow，如正在做 TPU 训练，应该考虑探索谷歌的 JAX。JAX 本身不是神经网络框架，更接近于具有自动微分能力的 GPU/TPU 的 NumPy 实现。DeepMind 的 Haiku，称为「Sonnet for JAX」，一个建立在 JAX 之上的神经网络库，如果正在考虑 JAX，Haiku 值得去探索。如果不进行 TPU 训练，最好是坚持使用 PyTorch。
无论选择哪种框架，都应该在 2022 年密切关注 JAX。
总之，在学术界，如果需要选择一个框架，请选择 PyTorch。
如果是一名教授怎么办？

 如果是一位教授，在深度学习课程中使用哪种框架取决于每一门课程的目标。如果课程的重点是培养具备行业技能的深度学习工程师，可以胜任整个端到端深度学习任务，不仅仅是掌握深度学习理论，应该使用 TensorFlow。在这种情况下，接触 TensorFlow 生态系统，工具以及端到端的实践项目，将非常有价值。

如果课程的重点是深度学习理论和理解深度学习模型的底层原理，应该使用 PyTorch。如果正在教授让学生为深度学习研究做好准备的高级本科课程或早期研究生课程，更应该选 PyTorch。
理想情况下，学生应该接触每个框架，尽管单个学期的时间有限，但花一些时间了解框架之间的差异可能很有价值。如果课程是机器学习大型课程的一部分，有许多课程专门针对不同的主题，最好坚持使用最适合课程材料的框架，而不是尝试同时接触两者。
如果正在尝试寻求职业的转变该怎么办？

 如果想在职业生涯中寻求一些转变，PyTorch 或 TensorFlow 都是不错的选择。在这种情况下，可以做的最重要的事情是证明可以带来额外价值，因此拥有复合项目经验至关重要。将深度学习应用于创造性用例来打破常规，或者通过端到端的项目表明已做好迎接行业的准备，都能处于不败之地。

因此，建议使用可以工作更轻松的任何框架，这比熟悉特定框架的 API 重要得多。如果对框架完全不熟悉，使用 TensorFlow，因为是首选的行业框架。下图中汇总了各种职业网站的每个框架职位的发布数量，TensorFlow 远远超过了 PyTorch。

 总之，如果有使用 PyTorch 的特定原因，如为 OpenAI 工作，或者使用 TensorFlow非常不直观，那么随意。但一般的建议是专注于 TensorFlow。

如果是业余爱好者，怎么选？

 如果是个深度学习爱好者，使用的框架将取决于目标。如果将深度学习模型作为某个较大项目的一部分来实施，TensorFlow 可能是想要使用的，尤其是在部署到物联网/嵌入式设备时。鉴于 PyTorch Live 已经发布，可以将 PyTorch 用于移动应用，但目前 TensorFlow + TFLite 仍然是首选方法。

如果目标是为了使用深度学习而学习，使用哪种框架取决于背景。一般来说，PyTorch 可能是更好的选择，特别是如果一直习惯使用 Python 工作。
如果是个初学者，怎么选？

 如果是一个对深度学习感兴趣，只是想入门的初学者，建议使用 Keras。利用高级组件，可以轻松地开始了解深度学习的基础知识。不过，一旦准备更深入地了解深度学习的具体细节，会面临几个选择：

如果不想安装新框架，担心能否适应新的 API，可以尝试从 Keras 「下降」到 TensorFlow。如果 TensorFlow困惑，尝试转向 PyTorch。
如果想要一个更像 Python 的框架，迁移到 PyTorch 可能是最好的选择。在这种情况下，必须安装新框架并可能重写自定义脚本。此外，如果 PyTorch有点麻烦，可以使用 PyTorch Lightning 划分代码，摆脱一些样板。
如果是一个完完全全的初学者，可以考虑观看一些 TensorFlow 和 PyTorch 的 YouTube 教程，以确定哪个框架更直观。
写在最后
PyTorch 与 TensorFlow 之争是一场微妙的辩论，格局一直在不断变化。至少在 2022 年，PyTorch 和 TensorFlow 都已经是非常成熟的框架，二者的核心深度学习功能重叠明显。今天，考虑每个框架的实际因素，如模型可用性、部署时间和相关的生态系统，都比只讨论技术差异更有意义。
选择任一框架都不会出错，因为都有完备的文档、学习资源和活跃的社区。希望能选到最适合的那一个框架。

参考链接

https://www.assemblyai.com/blog/pytorch-vs-tensorflow-in-2022/