pytorch/benchmark

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

PyTorch Benchmarks – What It Is

PyTorch Benchmarks is an open‑source suite that bundles a collection of popular deep‑learning models (e.g., BERT, ResNet, Stable Diffusion) and provides a standardised API so they can be run and timed consistently. The repo ships tiny versions of the data and scripts that install each model’s extra Python dependencies, letting you measure how fast a given PyTorch build (different versions, CUDA builds, TorchInductor, TorchScript, etc.) performs on a variety of workloads.


Why It Exists

  • Researchers and engineers need reproducible numbers when they change the PyTorch version, compiler flags, or hardware.
  • By exposing a common interface, the suite lets CI pipelines automatically run a set of “real‑world” models and compare results across nightly builds.

Getting Started (Installation)

  1. Set up Python – Python 3.8+ is supported; the README recommends a fresh Conda env with Python 3.11.
  2. Install CUDA libraries (if you have an NVIDIA GPU):
    conda install -y -c pytorch magma-cuda121   # CUDA 12.1 is default
    
  3. Install PyTorch, torchvision, torchaudio – either via Conda (nightly channel) or pip, but don’t mix the two package managers.
  4. Clone and install the benchmark suite:
    git clone https://github.com/pytorch/benchmark
    cd benchmark
    python3 install.py
    
    This script pulls in the per‑model dependencies and makes the torchbenchmark package importable.
  5. (Optional) Install as a library for use in other code:
    pip install git+https://github.com/pytorch/benchmark.git   # or `pip install .`
    

Running Benchmarks

Method What it does Typical command
test.py Quick sanity‑check – runs each model’s forward pass once. python3 test.py
test_bench.py Pytest‑benchmark driver – collects timing statistics, supports filtering, auto‑saves JSON results. pytest test_bench.py -k "test_BERT_pytorch_train_cpu"
run.py Simple CLI for a single model (debug/profiling). python3 run.py resnet50 -d cuda -t eval --profile
userbenchmark Framework for custom benchmark definitions; driven by run_benchmark.py. python run_benchmark.py my_custom_bench

Filters use the usual pytest -k expression syntax, letting you pick a model, device, or mode (train/eval).


Using the Models as a Library

You can import any model directly:

import torchbenchmark.models.densenet121 as densenet
bench = densenet.Model(test="eval", device="cuda", batch_size=1)
model, inputs = bench.get_module()
model(*inputs)   # runs a single forward pass

This is handy for CI tests or for embedding the models in your own profiling scripts.


Machine‑Specific Tuning

The suite includes utilities for low‑noise benchmarking on an AWS g4dn.metal instance (Amazon Linux). Running:

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

sets CPU governor, disables turbo, and records the configuration. The same script is invoked automatically when you run the pytest driver, but you can skip it with --ignore_machine_config.


CI Integration & Scores

Nightly CI runs the full model set against PyTorch nightly builds and publishes two score sets (V0 and V1) in the repository. Internally at Meta these scores feed a dashboard (Unidash), but the JSON outputs are also available publicly for anyone to analyse.


Extending the Suite

To add a new model, follow the Adding Models guide (torchbenchmark/models/ADDING_MODELS.md). You provide:

  • A thin wrapper that implements the standard BenchmarkModel API (get_module, train, eval, etc.).
  • Minimal data download script.
  • Optional support for alternative back‑ends (e.g., TorchInductor).

TL;DR

  • What: A curated set of deep‑learning workloads with a uniform interface for measuring PyTorch performance.
  • Who: PyTorch developers, hardware vendors, and anyone who needs reproducible speed numbers.
  • How: Install via Conda/pip, run test_bench.py (or run.py for a single model), and optionally tune the AWS g4dn.metal machine for low‑noise results.
  • Extendable: Add new models by implementing the small BenchmarkModel contract.

The repository is purely about benchmarking PyTorch; it does not provide training pipelines or novel models beyond the benchmark wrappers.

Related

  • Project
  • Project
  • Dispatch
  • Project
  • Project