TorchBench: one model API over rewritten PyTorch workloads

Every model rewritten behind one driver API

The workloads do not come from upstream. Inside torchbenchmark/models, each popular or exemplary model is copied and then modified in three ways: it exposes a standardised API that benchmark drivers call, it can optionally enable backends such as torchinductor or torchscript, and it carries a miniature version of its train and test data plus a dependency install script. A driver therefore never imports a model from its original repository. It imports the rewritten copy and asks it for a module and matching example inputs, which is what lets a single pytest filter expression select a model, an execution mode and a device at once. The API itself is small enough to read in one line: a model class is constructed with a test mode, a device and a batch size, and get_module() returns the module together with the inputs it expects. Each model directory also installs its own non-torch dependencies, on the assumption that the torch family is already present.

Torch, torchvision and torchaudio cannot be mixed

Installing torch is a separate job from installing the suite, deliberately, so that different torch versions can be benchmarked against each other. The constraint the project spells out is that all three packages must come from the same build process. A conda torchvision cannot be paired with a pip torch, and a torch compiled from source cannot be paired with a pip torchvision, because the packaging systems compile C extensions differently and produce incompatible binary extensions. Even the conda channel has to match: nightly against regular, never one of each. Where CUDA is involved, CUDA 11.8 and CUDA 12.1 are both supported with CUDA 12.1 as the default, and magma-cuda121 comes from the pytorch channel to match. The source-build route follows the same rule, with torchvision and torchaudio compiled locally too and USE_CUDA=1 set when CUDA is in play. Python 3.8 is the floor and 3.11 is what the project recommends, yet pyproject.toml still sets black to target-version py38, so the formatter is pinned to the oldest supported interpreter rather than the recommended one. Stable torch builds are not a target here: the suite asks for the newest nightly and says stable versions are not tested or maintained.

Three commands put the suite on a machine

Start from a Python 3.11 environment. Conda is optional but suggested.

# Using your current conda environment:
conda install -y python=3.11

# Or, using a new conda environment:
conda create -n torchbenchmark python=3.11
conda activate torchbenchmark

On an NVIDIA GPU machine, the matching CUDA build packages go in before torch itself:

conda install -y -c pytorch magma-cuda121

Then torch itself, from the nightly conda channels:

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

or from the nightly pip index, choosing one packaging system for the whole family:

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

Finally clone the repository and let the installer pull in per-model dependencies:

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

That last step is recursive across all models, so it needs working network access, and a proxy is the documented fix when the network requires one. The repository is intended to be installed from the source tree rather than from a pre-built wheel.

Automated tuning targets g4dn.metal and nothing else

Benchmark numbers move on a noisy machine, and this project treats that as a first-class problem rather than a footnote. Interrupts, context switches and clock frequency scaling conspire to make a run variable, so the suite ships utilities for automated tuning on one specific machine type: an AWS g4dn.metal instance running Amazon Linux. That instance family exposes processor state control and Amazon Linux has documented tuning guides, which is why the automation targets it and not another. Most of the same steps should be possible on Ubuntu, but nothing automates them there.

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

The command needs root because it changes system settings. Afterwards, pytest invokes machine_config on every run to assert the machine is still configured as expected and to write the configuration into the output json next to the measurements. Running pytest on an untuned machine is possible by passing --ignore_machine_config, and doing so drops those assertions as well as the recorded config metadata. The cost of this design is portability: anyone outside that single instance type has to reproduce the configuration by hand and gets no check that it stuck.

Three drivers, and the richest one is marked for replacement

Three entry points share the same infrastructure. test.py is the simplest wrapper, walking each model, installing it and executing its API as a sanity check rather than collecting timings. test_bench.py is the benchmarking driver, built on pytest-benchmark, and it collects benchmark statistics while accepting ordinary pytest filtering. userbenchmark exists for people writing their own benchmarks against TorchBench models.

python3 test.py -k "test_BERT_pytorch_train_cpu"

Generated test names follow a single pattern assembled from the model name, the execution mode and the device:

"test_" + <model_name> + "_" + {"train" | "eval" } + "_" + {"cpu" | "cuda"}

Because that naming is stable, pytest expressions can address one flavour of a model or a whole class of them:

-k "test_train[NAME-cuda]"

-k "(BERT and (not cuda))"

Handy options include --benchmark-autosave for json output, --collect-only to list what would run without running it, and --cpu_only for a local CPU box where the machine configuration checks do not apply. Two details deserve attention. The README first calls test_bench.py a pytest-benchmark script and then, a few lines later, says it is based on unittest, which are different runners, so the harness behind that driver is not settled by the text. It also states that test_bench.py will eventually be deprecated as the userbenchmark work evolves, so the driver with the widest option set is the one flagged for replacement.

Top-level scripts the README leaves unexplained

The repository root carries machinery the README never describes. bisection.py, regression_detector.py, run.py, run_benchmark.py, run_e2e.py and gen_summary_metadata.py sit next to test.py, test_bench.py and test_imports.py, joined by a docker/ directory, a .ci/ directory, a submodules/ directory, a .gitmodules file and CITATION.cff. What each script does, what triggers a bisection and where regression_detector.py reads its thresholds is written down nowhere in the README, which stops mid-sentence inside the userbenchmark section and never returns to them. setup.py is more telling: it declares the package as torchbench at version 0.1, includes only the torchbenchmark* and userbenchmark* packages through find_packages, and lists no install_requires at all. Dependencies live in requirements.txt instead, which pins timm==1.0.19 and transformers==4.57.3, requires nvidia-ml-py>=13.0.0 and numba>=0.57.0, and pulls in pytest-benchmark, submitit, boto3 and psutil. Those pins sit awkwardly beside the claim that the suite is self contained in terms of dependencies, except for the torch products, since timm and transformers are not torch products and install.py installs them recursively for every model.

One model import versus the whole driver

The README documents two narrower ways to use this code, and both undercut the full suite. The first is to skip the drivers and import one rewritten model as a library, after installing with pip install . or from the git URL:

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

That path returns a model and its inputs with no test runner, no machine_config assertion and no statistics, which is the right trade when one network is being profiled rather than builds compared. The second alternative is plain pytest-benchmark, which test_bench.py itself sits on: point it at any callable and it will time it, with no model installation, no train/eval and cpu/cuda naming scheme and no tuning utility. The suite earns its overhead when every model has to be measured the same way on a tuned machine; it is dead weight when a single callable is all that needs timing.

Editorial conclusion

pytorch/benchmark earns its place on any machine where torch, torchvision and torchaudio will be rebuilt or upgraded, because it fixes the model API, the test naming and the machine noise instead of leaving each to the person writing the script. It is the wrong tool for a quick look at one layer or one operator, where importing a single model beats installing the whole suite. Before trusting any number it prints, check that the three torch packages came from one build process, and read the run scripts the README leaves undocumented.

Frequently asked questions

Official sources

1. Issues
2. License: BSD-3-Clause
3. pytorch/benchmark on GitHub
4. README