在 Windows CPU 上提升 PyTorch 性能之路

作者：英特尔公司

PyTorch 在 Windows CPU 上的性能低于 Linux，这一直是一个重要问题。导致这种性能差异的因素有很多。通过我们的调查，我们找出了导致 Windows 上 CPU 性能低下的几个原因，其中两个主要问题已被确定：Windows 默认 malloc 内存分配器的效率低下以及 Windows 平台缺少用于向量化优化的 SIMD。在本文中，我们将展示 PyTorch CPU 在 Windows 上的性能如何较之前版本有所提升，以及在 PyTorch 2.4.1 版本中的表现。

PyTorch 2.1.2 及更高版本中的内存分配优化

在 PyTorch 2.1.2 之前的版本中，PyTorch 依赖操作系统的默认 malloc 函数进行内存分配。Windows 平台上的默认 malloc 内存分配机制比 Linux 平台上的 malloc 实现机制效率低，导致内存分配时间增加，性能下降。为了解决这个问题，我们将默认的 Windows malloc 替换为 mimalloc，这是一个由微软开发、更高效的内存分配器。如图 1.1 所示，这一更新已随 PyTorch 2.1.2 及更高版本的发布而包含，显著提升了 PyTorch 在 Windows 上的 CPU 性能。

通过内存分配优化提升 Windows 上 PyTorch CPU 性能

图 1.1：从 Windows PyTorch 2.0.1 版本升级到 2.1.2 版本实现的相对吞吐量提升（越高越好）。

该图表表明，随着 PyTorch 2.1.2 版本的发布，Windows 平台上的 CPU 性能有了显著提升。不同模型的改进程度各不相同，这归因于它们执行的操作组合的多样性及其相应的内存访问模式。虽然 BERT 模型显示出适度的性能提升，但像 ResNet50 和 MobileNet-v3 Large 这样的模型则受益于更显著的改进。

在高性能 CPU 上，内存分配成为性能瓶颈。这也是为什么解决此问题带来了如此显著的性能提升。

如下图所示，我们看到 PyTorch CPU 在 Windows 上的性能可以显著提升。然而，与在 Linux 上的性能相比，仍然存在明显的差距。PyTorch CPU 的 Windows 版本中缺少向量化优化是造成剩余性能差距的关键因素。

PyTorch 2.0.1 版本在 Windows 与 Linux 上的性能对比

图 1.2：PyTorch 2.0.1 版本在 Windows 与 Linux 上的相对性能（越高越好）。

PyTorch 2.1.2 版本在 Windows 与 Linux 上的性能对比

图 1.3：PyTorch 2.1.2 版本在 Windows 与 Linux 上的相对性能（越高越好）。

PyTorch 2.4.1 及更高版本中的向量化优化

在 PyTorch 2.4.1 之前，PyTorch 的 Windows 版本缺少用于向量化优化的 SIMD，而 Linux 版本利用了这一特性来提高性能。这种差异是由于 SLEEF 库在 Windows 上的集成问题，SLEEF 是一个用于评估基本函数的 SIMD 库，包括向量化 libm 和 DFT，对于高效的三角函数计算至关重要。通过与 ARM 和 Qualcomm 的工程师协作，这些挑战得到了解决，从而实现了 SIMD 在 Windows 上的 PyTorch 中的集成。因此，PyTorch 2.4.1 更新显著提升了 PyTorch 在 Windows 上的 CPU 性能，如图 2.1 所示。

通过向量化优化提升 Windows 上 PyTorch CPU 性能

图 2.1：从 PyTorch CPU 2.1.2 版本升级到 2.4.1 版本实现的相对吞吐量提升（越高越好）。

如下图所示，我们看到 PyTorch CPU 在 Windows 上的性能已达到在 Linux 上的性能水平。

PyTorch 2.4.1 版本在 Windows 与 Linux 上的性能对比

图 2.2：PyTorch 2.4.1 版本在 Windows 与 Linux 上的相对性能（越高越好）。

结论

从 PyTorch 2.0.1 到 PyTorch 2.4.1，Windows 与 Linux 之间的 CPU 性能差距持续缩小。我们比较了不同版本 PyTorch 在 Windows 上 CPU 性能与在 Linux 上 CPU 性能的比率，结果如下所示。

不同版本 PyTorch 在 Windows 与 Linux 上的性能对比

图 3：不同版本 PyTorch 在 Windows 与 Linux 上的性能比（越高越好）。

图表显示，使用 PyTorch 2.4.1，Windows 上的 CPU 性能已几乎与 Linux 上的性能趋同，在某些模型上甚至超越了 Linux。例如，对于 DistillBERT 和 RoBERTa 模型，Windows 与 Linux 的 CPU 性能比达到了惊人的 102%。然而，包括 MobileNet-v3 在内的某些模型仍存在性能差异。英特尔工程师将继续与 Meta 工程师协作，以缩小 PyTorch CPU 在 Windows 与 Linux 之间的性能差距。

如何利用这些优化

从官方仓库在 Windows 上安装 PyTorch CPU 2.4.1 或更高版本，即可自动体验内存分配和向量化带来的性能提升。

致谢

本博客文章中呈现的结果是英特尔 PyTorch 团队与 Meta 协作努力的成果。我们要衷心感谢 Xu Han、Jiong Gong、Haozhe Zhu、Mingfei Ma、Chuanqi Wang、Guobing Chen 和 Eikan Wang。他们在实现本文讨论的优化和性能改进方面发挥了重要作用。感谢来自社区的 Jiachen Pu 参与问题讨论并建议使用 mimalloc。我们还要感谢微软提供了如此易于集成且高性能的内存分配库。感谢来自 ARM 的 Pierre Blanchard、Nathan Sircombe 以及来自 Adobe 的 Alex Reinking 在克服 sleef 集成到 PyTorch Windows 中的兼容性问题方面做出的贡献。最后，我们感谢 Jing Xu、Weizhuo Zhang 和 Zhaoqiong Zheng 对这篇博客的贡献。

产品与性能信息

表中配置信息通过 svr-info 收集。由英特尔于 2024 年 8 月 30 日测试。

注意事项和免责声明

性能因使用方式、配置和其他因素而异。在性能指标网站上了解更多信息。

性能结果基于所示日期的测试配置，可能不反映所有公开可用的更新。配置详情请参见备份。任何产品或组件都无法做到绝对安全。您的成本和结果可能会有所不同。英特尔技术可能需要启用相应的硬件、软件或服务激活。

英特尔公司。英特尔、英特尔徽标和其他英特尔标记是英特尔公司或其子公司的商标。其他名称和品牌可能属于他人财产。