pytorch/benchmark

TorchBench is a collection of open source benchmarks used to evaluate PyTorch performance.

PyTorch Benchmarks – 它是什么

PyTorch Benchmarks 是一个开源工具集，包含一系列流行深度学习模型（例如 BERT、ResNet、Stable Diffusion）并提供一个标准化 API，使它们能够以一致的方式运行和计时。该仓库包含每个模型的极小数据集和安装其额外 Python 依赖项的脚本，让你可以测量特定 PyTorch 构建（不同版本、CUDA 构建、TorchInductor、TorchScript 等）在各种工作负载上的性能表现。

为什么存在

研究人员和工程师在更改 PyTorch 版本、编译器标志或硬件时需要可复现的性能数据。
通过暴露统一接口，该工具集允许 CI 流水线自动运行一组“真实世界”模型，并在夜间构建之间进行结果对比。

快速开始（安装）

设置 Python – 支持 Python 3.8+；README 推荐使用 Python 3.11 的全新 Conda 环境。

安装 CUDA 库（如果你有 NVIDIA GPU）：

conda install -y -c pytorch magma-cuda121   # CUDA 12.1 为默认值

安装 PyTorch、torchvision、torchaudio – 可通过 Conda（nightly 频道）或 pip 安装，但不要混用两个包管理器。
克隆并安装基准测试套件：
```
git clone https://github.com/pytorch/benchmark
cd benchmark
python3 install.py
```
此脚本会拉取各模型的依赖项，并使 torchbenchmark 包可导入。

（可选）作为库安装，以便在其他代码中使用：

pip install git+https://github.com/pytorch/benchmark.git   # 或 `pip install .`

运行基准测试

方法	作用	常用命令
`test.py`	快速健康检查 – 每个模型运行一次前向传播。	`python3 test.py`
`test_bench.py`	Pytest-benchmark 驱动器 – 收集时间统计信息，支持过滤，自动保存 JSON 结果。	`pytest test_bench.py -k "test_BERT_pytorch_train_cpu"`
`run.py`	单模型的简单 CLI（用于调试/性能分析）。	`python3 run.py resnet50 -d cuda -t eval --profile`
`userbenchmark`	自定义基准定义框架；由 `run_benchmark.py` 驱动。	`python run_benchmark.py my_custom_bench`

过滤器使用标准的 pytest -k 表达式语法，可选择模型、设备或模式（train/eval）。

将模型作为库使用

你可以直接导入任意模型：

import torchbenchmark.models.densenet121 as densenet
bench = densenet.Model(test="eval", device="cuda", batch_size=1)
model, inputs = bench.get_module()
model(*inputs)   # 执行一次前向传播

这在 CI 测试或嵌入到自己的性能分析脚本中非常方便。

机器特定调优

该套件包含用于在 AWS g4dn.metal 实例（Amazon Linux）上进行低噪声基准测试的实用工具。运行：

sudo $(which python) torchbenchmark/util/machine_config.py --configure

将设置 CPU 调度器、禁用 Turbo 模式，并记录配置。当运行 pytest 驱动器时，该脚本会自动调用，但你可以通过 --ignore_machine_config 跳过它。

CI 集成与分数

夜间 CI 会针对 PyTorch 夜间构建运行全部模型集，并在仓库中发布两组分数（V0 和 V1）。Meta 内部这些分数会输入仪表板（Unidash），但 JSON 输出也对所有人公开，可供分析。

扩展套件

要添加新模型，请遵循 添加模型 指南（torchbenchmark/models/ADDING_MODELS.md）。你需要提供：

实现标准 BenchmarkModel API（get_module、train、eval 等）的轻量封装。
最小化数据下载脚本。
可选地支持替代后端（例如 TorchInductor）。

TL;DR

是什么：一组经过筛选的深度学习工作负载，具有统一接口，用于测量 PyTorch 性能。
谁适用：PyTorch 开发者、硬件厂商，以及需要可复现速度数据的任何人。
如何使用：通过 Conda/pip 安装，运行 test_bench.py（或单模型用 run.py），并可选地在 AWS g4dn.metal 机器上进行低噪声结果调优。
可扩展：通过实现小的 BenchmarkModel 合约来添加新模型。

该仓库仅专注于 PyTorch 基准测试；它不提供训练流水线或超出基准封装的新模型。

pytorch/benchmark