PyTorch 基准测试

这是一个用于评估 PyTorch 性能的开源基准测试集合。

torchbenchmark/models 包含一些流行或典型的负载副本，这些副本已被修改，以： (a) 暴露一个标准化的 API 供基准测试驱动程序使用；(b) 可选地支持诸如 torchinductor/torchscript 等后端；(c) 包含训练/测试数据的精简版本以及依赖项安装脚本。

安装

该基准测试套件在依赖项方面应是自包含的， 但 PyTorch 相关产品除外，因为它们需要单独安装，以便可以对不同版本的 PyTorch 进行基准测试。

使用预构建包

我们支持 Python 3.8 及以上版本，推荐使用 3.11。Conda 是可选的，但建议使用。要在 Conda 中使用 Python 3.11：

# 使用当前的 Conda 环境：
conda install -y python=3.11

# 或者，使用一个新的 Conda 环境：
conda create -n torchbenchmark python=3.11
conda activate torchbenchmark

如果您要运行 NVIDIA GPU 测试，我们同时支持 CUDA 11.8 和 12.1，并默认使用 CUDA 12.1：

conda install -y -c pytorch magma-cuda121

然后使用 Conda 安装 PyTorch、torchvision 和 torchaudio：

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch-nightly -c nvidia

或者使用 pip： （但请勿将 pip 和 Conda 混用来安装 PyTorch 家族的库！— 参见下方说明）

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121

安装基准测试套件，它会递归地安装所有模型所需的依赖项。目前，该仓库建议从源代码树进行安装。

git clone https://github.com/pytorch/benchmark
cd benchmark
python3 install.py

将 torchbench 作为库安装

如果您希望将 torchbench 作为库来运行，可以执行以下操作：

python3 install.py
pip install git+https://www.github.com/pytorch/benchmark.git

或者

python3 install.py
pip install . # 添加 -e 选项以进行可编辑安装

上述操作后，您可以这样使用：

import torchbenchmark.models.densenet121
model, example_inputs = torchbenchmark.models.densenet121.Model(test="eval", device="cuda", batch_size=1).get_module()
model(*example_inputs)

从源码构建

请注意，当从源码构建 PyTorch 时，torchvision 和 torchaudio 也必须从源码构建，以确保 C API 匹配。

有关安装 torchvision 的详细说明，请参阅 此处，而 torchaudio 的说明则可在 此处找到。如果使用 CUDA，请务必启用 CUDA（通过 USE_CUDA=1）。之后，执行以下命令：

git clone https://github.com/pytorch/benchmark
cd benchmark
python3 install.py

注意事项

• 设置步骤需要网络连接——如有必要，请确保已启用代理。
• 我们建议使用最新的 PyTorch 夜间版来运行基准测试。稳定版未经过测试或维护。
• PyTorch、torchvision 和 torchaudio 必须来自同一构建过程。这意味着不能将 Conda 版本的 torchvision 与 pip 版本的 PyTorch 混用，也不能将源码编译的 PyTorch 与 pip 安装的 torchvision 混用。即使是 Conda 通道（夜间版 vs 正常版）也必须一致。这是由于不同打包系统采用的编译流程存在差异，从而导致生成的 Python 二进制扩展不兼容。

使用低噪声机器

各种噪声源，如中断、上下文切换、时钟频率缩放等，都可能导致基准测试结果出现波动。在根据基准测试数据得出结论之前，了解您环境中的噪声水平非常重要。虽然原则上任何机器都可以进行优化，但具体步骤和最终效果会因操作系统、内核、驱动程序和硬件的不同而有所差异。为此，torchbenchmark 选择了一种能够良好支持的机器类型，并提供了针对该机器的自动化调优工具。未来，我们可能会支持更多类型的机器，并欢迎在此方面的贡献。

目前支持的机器类型是使用 Amazon Linux 的 AWS g4dn.metal 实例。这是 AWS 实例类型中少数支持 处理器状态控制的实例之一，并且有针对 Amazon Linux 的官方调优指南。尽管这些步骤在大多数情况下也可以在 Ubuntu 上实现，但尚未实现自动化。

要调优您的 g4dn.metal Amazon Linux 机器，运行以下命令：

sudo `which python` torchbenchmark/util/machine_config.py --configure

在运行 pytest 时（见下文），machine_config 脚本会被调用以确认配置是否正确，并将配置信息记录到输出的 JSON 文件中。如果希望在不进行调优的情况下运行 pytest，可以使用 --ignore_machine_config 选项。

运行模型基准测试

运行模型基准测试有多种方式。

test.py 提供了一个最简单的封装，用于遍历每个模型并安装和执行它们。

test_bench.py 是一个 pytest-benchmark 脚本，它利用相同的基础设施，但会收集基准测试统计信息，并支持 pytest 过滤功能。

userbenchmark 允许开发和运行自定义基准测试。

在每个模型仓库中，假设用户已经安装了 PyTorch 家族的所有软件包（PyTorch、torchvision、torchaudio 等），但它会安装该模型所需的其余依赖项。

使用 test.py

python3 test.py 会执行每个模型的 API，作为一项健全性检查。若要进行基准测试，请使用 test_bench.py。它基于 unittest，并支持通过 CLI 进行过滤。

例如，要在 CPU 上运行 BERT 模型的训练模式：

python3 test.py -k "test_BERT_pytorch_train_cpu"

测试名称遵循以下模式：

"test_" + <模型名称> + "_" + {"train" | "eval"} + "_" + {"cpu" | "cuda"}

使用 pytest-benchmark 驱动程序

pytest test_bench.py 会调用基准测试驱动程序。完整的选项列表可通过 --help 查看。

一些有用的选项包括：

• --benchmark-autosave（或其他保存相关标志）以获取 .json 格式的输出
• -k <筛选表达式> 标准的 pytest 筛选功能
• --collect-only 仅显示将要运行的测试，可用于查看有哪些模型或调试您的筛选表达式
• --cpu_only 如果在本地 CPU 机器上运行且忽略机器配置检查

基准测试筛选示例

• -k "test_train[NAME-cuda]" 用于特定模型的特定变体
• -k "(BERT and (not cuda))" 用于更灵活的筛选方法

请注意，随着 userbenchmark 的不断发展，test_bench.py 最终将被弃用。鼓励用户探索并考虑使用 userbenchmark。

使用 userbenchmark

userbenchmark 允许您基于 TorchBench 中的模型开发自定义基准测试。请参阅 userbenchmark 使用说明，了解如何创建新的 userbenchmark。随后，您可以使用 run_benchmark.py 驱动程序来运行基准测试，例如：python run_benchmark.py <benchmark_name>。运行 python run_benchmark.py —help 可查看可用选项。

使用 run.py 进行简单调试或性能分析

有时您可能只想针对特定模型运行训练或评估，例如用于调试或性能分析。为了避免依赖每个模型内部的 __main__ 实现，run.py 提供了一个轻量级的命令行界面来实现这一目的，该界面基于标准的 BenchmarkModel API 构建。

python3 run.py <model> [-d {cpu,cuda}] [-t {eval,train}] [--profile]

注意： 可以是完整的精确名称，也可以是部分字符串匹配。

将 torchbench 模型用作库

如果您希望将 torchbench 作为一个可供测试的模型集合来使用，最简单的方式是将其集成到您的代码、CI 或测试中，如下所示：

import torch
import importlib
import sys

# 如果您的目录结构为 this_file.py, benchmark/
sys.path.append("benchmark")
model_name = "torchbenchmark.models.stable_diffusion_text_encoder" # 替换为您正在使用的模型名称
module = importlib.import_module(model_name)

benchmark_cls = getattr(module, "Model", None)
benchmark = benchmark_cls(test="eval", device="cuda") # test 可为 train 或 eval；device 可为 cuda 或 cpu

model, example = benchmark.get_module()
model(*example)

每日 CI 测试运行

目前，这些模型在每日构建的 PyTorch 版本上运行，并将数据推送至 Meta 的内部数据库。每日 CI 会发布 V1 和 V0 两种版本的性能分数。

详情请参见 Unidash（仅限 Meta 内部）。

添加新模型

请参阅 添加模型。

PyTorch Benchmark 快速上手指南

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统：Linux (推荐 Amazon Linux 或 Ubuntu)，Windows/macOS 支持有限。
• Python 版本：Python 3.8+ (推荐使用 Python 3.11)。
• GPU 支持：若需测试 NVIDIA GPU，需安装 CUDA 11.8 或 12.1 (默认推荐 CUDA 12.1)。
• 网络环境：安装过程需要下载大量依赖，请确保网络畅通。国内用户建议配置代理或使用国内镜像源加速 pip/conda。
• 管理工具：推荐使用 Conda 管理环境（可选但强烈建议）。

注意：torch, torchvision, torchaudio 必须来自同一构建源（例如全部通过 conda 安装，或全部通过 pip 安装，或全部源码编译），严禁混用，否则会导致二进制不兼容。

安装步骤

1. 创建并激活 Conda 环境

推荐使用 Python 3.11 创建独立环境：

# 创建新环境
conda create -n torchbenchmark python=3.11
conda activate torchbenchmark

2. 安装 PyTorch 及相关组件

根据是否使用 GPU 选择以下一种方式安装。国内用户建议使用清华或中科大镜像源加速。

方案 A：使用 Conda 安装 (推荐)

# 安装 magma (仅 NVIDIA GPU 需要)
conda install -y -c pytorch magma-cuda121

# 安装 PyTorch 全家桶 (使用 nightly 版本以获得最新支持)
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch-nightly -c nvidia

方案 B：使用 Pip 安装

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121

3. 安装 Benchmark 套件

克隆仓库并运行安装脚本，该脚本会自动递归安装所有模型所需的依赖。

git clone https://github.com/pytorch/benchmark
cd benchmark
python3 install.py

提示：如果您希望将 torchbench 作为库在其他项目中使用，也可执行 pip install . 或 pip install -e . 进行安装。

基本使用

安装完成后，您可以通过以下几种方式使用该工具：

1. 运行特定模型测试 (最简单)

使用 test.py 可以快速验证单个模型的运行情况。命令格式遵循 test_<模型名>_<模式>_<设备> 的模式。

示例：在 CPU 上运行 BERT 模型的训练模式

python3 test.py -k "test_BERT_pytorch_train_cpu"

示例：在 CUDA 上运行 BERT 模型的评估模式

python3 test.py -k "test_BERT_pytorch_eval_cuda"

2. 执行性能基准测试

使用 pytest 驱动 test_bench.py 可以收集详细的性能统计数据。

示例：运行所有可用的基准测试并保存结果为 JSON

pytest test_bench.py --benchmark-autosave

示例：仅运行包含 "BERT" 且不在 CUDA 上运行的测试

pytest test_bench.py -k "(BERT and (not cuda))"

3. 调试与性能分析 (Profile)

如果只需对单个模型进行简单的调试或性能剖析，可使用 run.py。

示例：对 resnet18 模型在 CUDA 上进行训练模式剖析

python3 run.py resnet18 -d cuda -t train --profile

4. 作为 Python 库调用

您可以在自己的代码中直接加载模型进行测试：

import torchbenchmark.models.densenet121

# 初始化模型 (模式：eval, 设备：cuda, 批次大小：1)
model, example_inputs = torchbenchmark.models.densenet121.Model(test="eval", device="cuda", batch_size=1).get_module()

# 执行推理
model(*example_inputs)