PyTorch 基准测试

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

torchbenchmark/models 包含了经过修改的流行或exemplary工作负载的副本，这些修改旨在： (a) 为基准测试驱动程序提供标准化的API，(b) 可选地启用 torchinductor/torchscript 等后端， (c) 包含训练/测试数据的微型版本和依赖项安装脚本。

安装

除了 torch 产品（intended to be installed separately so different torch versions can be benchmarked）之外，基准测试套件在依赖项方面应该是自包含的。

使用预构建包

我们支持 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 安装 torch 系列库！- 见下方注释）

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 每日构建版本运行基准测试。稳定版本未经测试或维护。
torch、torchvision 和 torchaudio 必须都从同一构建过程安装。这意味着不可能混用 conda 安装的 torchvision 和 pip 安装的 torch，或混用从源代码构建的 torch 和 pip 安装的 torchvision。重要的是要匹配甚至 conda 频道（nightly 与 regular）。这是由于不同打包系统使用的编译过程产生了不兼容的 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允许开发和运行自定义基准测试。

在每个模型仓库中，假设用户已经安装了所有torch系列包(torch、torchvision、torchaudio等)，但它会安装模型所需的其他依赖项。

使用`test.py`

python3 test.py将执行每个模型的API，作为健全性检查。对于基准测试，请使用test_bench.py。它基于unittest，支持通过命令行进行过滤。

例如，要在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 <基准测试名称>。运行python run_benchmark.py --help查看可用选项。

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

有时您可能只想对特定模型运行训练或评估，例如用于调试或分析。与其依赖每个模型内部的__main__实现，run.py为此目的提供了一个轻量级CLI，基于标准BenchmarkModel API构建。

python3 run.py <模型> [-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)