PyTorch 与 TensorFlow 怎么选？

2022 年 8 月 15 日
北京
本文字数：2857 字
阅读完需：约 9 分钟

PyTorch 和 TensorFlow 是目前最主流的两个深度学习框架，绝大多数研究者会选择 PyTorch 或者 TensorFlow 进行深度学习的入门学习。图 1 展示了近两年来几个主流深度学习框架的 Google 指数，其中 PyTorch 和 TensorFlow 的热度不相上下，均遥遥领先于其他框架。

图 1 TensorFlow、PyTorch、MXNet、Caffe 和 PaddlePaddle 的 Google 指数

本文将从以下 4 个方面对比介绍 TensorFlow 和 PyTorch，帮助读者更好地选择学习和使用哪个框架。

1.动态图与静态图

几乎所有的框架都是基于计算图的，计算图可以分为静态图和动态图两种。静态图是先定义再运行（define and run），一次定义、多次运行，这也意味着一旦创建就不能修改。静态图在定义时使用了特殊的语法，就像学习一门新的语言，同时在构建静态图时需要考虑所有的情况，这导致静态图过于庞大，可能占用过多的显存，不过其运行速度相对较快。动态图是在运行过程中被定义的，即在运行时构建（define by run），可以多次构建，多次运行。在构建动态图时可以使用 Python 的 if、while、for-loop 等常见语句，构建更加简单、直观，同时可以方便地进行修改、调试。

TensorFlow 最初选择使用静态图，这样的设计带来了较高的性能，但在构建网络时较为烦琐，用户需要专门学习 TensorFlow 的语法架构才能搭建网络，同时很难调试。PyTorch 选择使用动态图，动态图的设计模式更加符合人类的思考过程，方便查看、修改中间变量的值，用户可以轻松地搭建网络进行训练。目前，TensorFlow 2.0 之后的版本已经支持动态图的构建，并且提供动态图与静态图的转换功能。

2.学术研究和开源代码

人工智能作为一个新兴学科，仍需进行大量研究、探索以解决各种问题。图 2 展示了近两年来几个主流深度学习框架在开源的学术论文中所占的比例（数据来源于 Paperswithcode），其中 PyTorch 凭借其易用性和大量的开源资源在学术界遥遥领先。以 2021 年 6 月的数据为例，PyTorch 所占的比例几乎是 TensorFlow 的 6 倍。随着开源资源数量的逐步增多，这种领先优势也会逐步扩大。

图 2 几个主流深度学习框架在开源的学术论文中所占的比例

3.工业化

工业界更加关心如何将深度学习算法部署到各种架构的平台上，TensorFlow 凭借其完备的功能支持受到了工业界的广泛欢迎。TensorFlow 提供了 TensorFlow Serving 和 TensorFlow Lite，可以便捷地将训练好的模型部署到集群以及移动设备上。同时，TensorFlow 开源较早，许多公司已经建立起完整的使用 TensorFlow 开发、部署的模式，这对于追求稳定性的工业界十分重要。不过，受限于 TensorFlow 2.0 版本后的接口变动，许多成熟的模型并不能直接在新版本的 TensorFlow 上运行。而 PyTorch 曾经在工业化部署方面相对较弱，但是背靠 Facebook 数十亿用户，近些年来，PyTorch 团队也在着手进行完善。PyTorch 于 2020 年 4 月发布了 TorchServe，以帮助用户灵活地进行模型部署。同时，PyTorch 对 ONNX（微软定义的一种开放式的文件格式）和 TensorRT 的支持也愈加丰富。然而，对于部分复杂的算子，仍然需要重新进行设计。

4.学习成本

TensorFlow 作为一个相对底层的系统，它创造了许多全新的概念，用户近似需要学习一门新的语言。TensorFlow 最令人诟病的是它混乱且频繁更新的接口设计：同样的功能提供了多种实现，对接口的设计没有考虑向后兼容性等。虽然有关 TensorFlow 的文档、教程很多，但是整体而言层次性不强，用户很难快速地使用 TensorFlow 完成具体的任务。

而 PyTorch 更多地从算法工程师的角度考虑，它的接口风格与 NumPy 类似，Python 用户可以便捷地使用 PyTorch 搭建模型，并进行调试。同时，PyTorch 提供了完整的文档、循序渐进的指南以及作者亲自维护的论坛供用户交流问题。

总体而言，TensorFlow 是一个十分完备的系统，在开发它时主要从系统设计的角度出发，目标十分宏大，力求成为最强大的深度学习框架。然而，作为一个还在快速发展探索的领域，人工智能研发人员实际需要的可能并没有这么复杂，他们更加希望快速地实现和尝试自己的想法，这也是为什么半路出家的 PyTorch 能够异军突起，一举成为最受欢迎的深度学习框架之一（可能没有“之一”）

为什么选择 PyTorch

PyTorch 作为最受欢迎的深度学习框架之一，它主要有以下几个核心优势。

• 易用：编程是一项智力劳动，通过易用的工具可以最快地实现用户的想法。最易用的程序就是用户需要学习新东西最少的程序，换句话说，最易用的程序就是最切合用户已有知识的程序。PyTorch 具有 Pythonic 设计风格，以及与 NumPy 类似的接口，Python 用户可以快速、方便地使用 PyTorch 搭建模型。同时，PyTorch 官方提供了层次分明的文档，用户的学习成本较低。基于动态图的优势，用户可以方便地调试 PyTorch 代码。PyTorch 的设计符合人们的思维，它让用户尽可能地专注于实现自己的想法，所思即所得，不需要考虑太多框架本身的束缚。

• 简洁：PyTorch 的设计追求更少的封装，尽量避免重复造轮子。PyTorch 的设计遵循 Tensor→autograd→nn.Module 三个由低到高的抽象层次，其分别代表高维数组、自动求导和神经网络。这三个抽象之间的关系紧密，可以同时进行修改和操作。简洁的设计带来的另一个好处就是代码易于理解。更少的抽象、更加直观的设计，使得 PyTorch 的源码十分容易阅读，使用 PyTorch 搭建的模型也更加清晰、直观。PyTorch 的设计真正做到了“Keep it Simple, Stupid”，与 UNIX 的简洁之美如出一辙。

• 生态：PyTorch 拥有最多的开源模型，这意味着无论研究什么领域，用户都可以轻松找到相关源码进行研读。关于深度学习的绝大多数前沿研究是使用 PyTorch 进行的，许多顶级会议开放的源码也使用 PyTorch 构建。如果读者想要学习最新的算法，并将其应用在实践中，那么 PyTorch 将是不二之选。同时，PyTorch 有着优秀的社区，用户可以方便地交流和求教问题。背靠 Facebook，PyTorch 的开发者也会根据用户的反馈不断迭代更新 PyTorch。

• 拓展：PyTorch 凭借其易用性与简洁性成为最受欢迎的深度学习框架之一。目前，许多国产的深度学习框架（如 Jittor、MegEngine 等）的接口风格均与 PyTorch 保持一致，以此降低用户的学习成本。读者在掌握了 PyTorch 的基本用法后，可以很容易学习和使用新的框架。

总体而言，PyTorch 十分适合作为用户学习的第一个深度学习框架，相信用户在学习和使用的过程中会逐渐爱上 PyTorch，并使用它打开深度学习和人工智能的大门。

本文节选自《深度学习框架 PyTorch：入门与实践（第 2 版）》一书，本书第 1 版于 2018 年出版，当时主流的 PyTorch 版本为 0.3。随着 PyTorch 的不断更新迭代，许多函数接口已经被舍弃，同时新增了许多更加实用的功能接口。在本书第 1 版出版后，许多读者通过各路渠道提出了不少好的建议与意见，综合这些建议与意见，作者对本书内容进行了较大程度的修订重新再版。