PyTorch VS TensorFlow谁最强？这是星标15000+ Transformers库的运行结果

关注

转载最新推荐文章于 2026-05-19 00:19:47 发布 600 阅读 AI 写同款· GEO 优化 ›

TensorFlow和Pytorch两种机器学习框架的比较及优缺点 TensorFlow是一个开源的机器学习框架，由Google开发并发布，它被用来构建各种类型的机器学习模型，例如图像识别、语音识别、自然语言处理等。自然语言处理：使用TensorFlow可以构建深度学习模型，例如循环神经网络（RNN）和长短期记忆网络（LSTM），用于自然语言处理任务，例如文本生成和情感分析等。选择哪一种框架，还需要根据具体的应用场景进行考虑。1. 易用性和灵活性带来的缺点是，PyTorch在大型数据集上需要额外的工作来优化它的计算效率，而且同时也影响了代码的可维护性。阅读详情

作者 | Lysandre Debut

译者 | 陆离

出品 | AI科技大本营（ID: rgznai100）

PyTorch VS TensorFlow谁最强？

这是标星15000+ Transformers库的运行结果

【导语】自然语言处理预训练模型库 Transformers 实现了几种用于 NLP 任务的最先进的 Transformer 架构，如文本分类、信息提取、问题解答和文本生成等，它经常被研究人员和公司所使用，提供 PyTorch 和 TensorFlow 的前端实现。

扫码加小编微信拉你入机器学习算法群，一起交流进步；

究竟是 PyTorch 还是 TensorFlow 更有效率地训练和运行 Transformers 模型？作者对不同环境下所展现的性能进行了对比，最终的结果是，无论在 CPU 还是 GPU 上，最终两大框架的表现都差不多。

Transformers库：

https://github.com/huggingface/transformers

自 TensorFlow 发布以来，我们一直在致力于模型产品化的工作，并使其可以用在 TPU 上，逐步适应它的性能。

本文对比了我们的模型在几种环境中所展现出来的性能。在 CPU 和 GPU 上比较了 PyTorch（1.3.0）和 TensorFlow（2.0）的推断结果。出于对一些原因的考虑，本文只是关于基准测试和后续性能优化系列文章中的第一篇。此外，我们还在文档中创建了基准测试部分，随着进一步的模型研究，并在不同环境中对它们进行基准测试，该部分还将不断完善。

结果

测试的平均结果显示在下表中，在接下来的讨论部分里将对这些结果进行详细介绍。

平均推断时间

Benchmarking Transformers测试结果：https://url.cn/5hZHCll

电子表格中的 N/A 记录表示内存不足或是数列长度不合适。Transformer-XL（Transformer extra-long，Transformer 在模型设计上做了长度方面的延申）没有 TorchScript 结果，因为它当前不能由 TorchScript（TorchScript 是一种从 PyTorch 代码创建可序列化和可优化模型的方法）进行序列化。

大多数情况下，TensorFlow 和 PyTorch 的模型在 GPU 和 CPU 上都得到了非常相似的结果。下面是对结果相关的阐述，不仅是 PyTorch 和 TensorFlow 之间的比较，也是模型之间的比较。

测量推理

推理时间是模型投入生产时的一个重要指标。为了评估模型的推理时间，我们对不同批量和不同序列长度的模型进行了对比。我们比较了适当的批量大小[1，2，4，8]和序列长度[8，64，128，256，512，1024]。批量大小仍然很小，因为我们仅仅关注推断设置。BERT和其它类似模型的最大序列长度为 512 或 256（用于 CTRL），因此不能在最后一个序列长度上进行测量。

以下是我们在两种不同环境中的测试结果：

在 CPU 上，使用 GCP n1-standard-32，它有 32 个 vCPU 和 120GB 的内存。CPU 的型号是 Intel Xeon@2.3GHz：

在 GPU 上，使用了带有 12 个 vCPU、40GB 内存和一个 V100 的 GPU（16GB 的 VRAM）的定制 GCP 机器；

实验细节和最佳实践

为了最大化性能，我们进行了更进一步的优化：

上述测量使用的 Intel Xeon CPU 带有 AVX 和 AVX2 的扩展，而 TensorFlow 需要从源代码编译之后才能够利用这些扩展，所以只能这么做；

我们通过使用 tf.function 和预先跟踪模型来确保我们没有使用 TensorFlow 的 eager 模式；

我们比较了依赖于库和不依赖于库的工具：PyTorch的 TorchScript 和带有 GPU 的TensorFlow 的 XLA（自动聚类），后面会详细介绍这两个工具；

我们使用了原生的 Python 模块 timeit 来测量推断时间。另外，用repeat=30 和 number=3 来进行每个实验。然后求出30 个值的平均数，就会得到所期望的平均推理时间。通常超过30 个值就会获得非常稳定的结果了；

我们不会使用如 TFX 这样的生产环境，并且我们使用的测量模型的可调用方法是：PyTorch 的nn.module.forward 和 TensorFlow 的tf.keras.layers.layer.call；

对于 TensorFlow 和 PyTorch，我们会很谨慎地使用适当的 CUDA 版本；

讨论

PyTorch 和 TensorFlow

在大多数情况下，这两个框架都会得到类似的结果，与 PyTorch 相比，TensorFlow 在CPU 上的速度通常会稍慢一些，而在 GPU 上的速度则稍快一点：

所有的模型中，在 CPU 上，PyTorch 的平均推断时间为 0.748s，而 TensorFlow 的平均推断时间为 0.823s；

所有模型中，在 GPU 上，PyTorch 的平均推断时间为 0.046s，而 TensorFlow 的平均推断时间为 0.043s；

这些结果通过计算平均值来比较所有模型的推断时间。因此，输入值越大，对最终结果的影响就越大。当输入值过大时，PyTorch 就会耗尽内存；当计算平均值时，这些结果会从所有度量中删除，因为这样会使结果向 PyTorch 倾斜。

运行过程中，PyTorch 模型往往比 TensorFlow 模型更早地耗尽内存：除了Distilled 模型之外，PyTorch 在输入批量大小达到 8 以及序列长度达到 1024 时会耗尽内存。

TorchScript

TorchScript 是PyTorch 用来创建可序列化模型的一种方法，可以在不同的运行时间上运行，而不需要 Python 的依赖包，如 C++ 环境。我们的测试是通过在Python 中跟踪模型并在相同的环境中重用这个跟踪模型来完成的。我们通过预先执行前向传递以确保在测量其推断之前来跟踪模型。

免责声明：虽然 TorchScript 并不是为了在 Python 环境中提高运行速度而创建的，但是我们的结果表明，使用 TorchScript 的跟踪模型可以提高性能。

TorchScript 似乎非常依赖于模型以及输入的大小（批量大小*序列长度）。例如，在XLNet 上使用 TorchScript 可以获得永久性的性能提升，而在 XLM 上使用 TorchScript 则可能会出现问题，因为它在较小输入下才会提高性能，但在较大输入下反而会降低性能。

平均来说，使用 TorchScript 跟踪模型的推断速度比使用相同 PyTorch 非跟踪模型的推断速度快 20%。