[torchbench] update environment setup script (#156465)

Existing torchbench `Makefile` installs all models from torchbench, which could easily take 30 minutes, even if a developer only want to run 1 model.

This PR adds a config to only install torchbench models we want to run.

Example usage:
```
# Install 1 torchbench model
make build-deps TORCHBENCH_MODELS="alexnet"

# Install 3 torchbench models
make build-deps TORCHBENCH_MODELS="alexnet basic_gnn_gcn BERT_pytorch"

# Install all models
make build-deps

# Install all models
make build-deps TORCHBENCH_MODELS=""
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/156465
Approved by: https://github.com/ezyang
This commit is contained in:
Boyuan Feng 2025-06-23 17:41:24 +00:00 committed by PyTorch MergeBot
parent e583b88819
commit a95504b10f

View File

@ -1,3 +1,9 @@
# Usage:
# make build-deps TORCHBENCH_MODELS=<model_names>
# Support install a single torchbench model (e.g., "alexnet"),
# or multiple torchbench model names (e.g., "alexnet basic_gnn_gcn BERT_pytorch"),
# or empty (i.e., "") for installing all torchbench models.
clone-deps:
(cd ../../.. \
&& (test -e torchvision || git clone --recursive https://github.com/pytorch/vision torchvision) \
@ -23,11 +29,11 @@ pull-deps: clone-deps
build-deps: clone-deps
uv pip install astunparse numpy scipy ninja pyyaml mkl mkl-include setuptools cmake \
typing-extensions requests protobuf numba cython scikit-learn librosa
(cd ../../../torchvision && uv pip install -e .)
(cd ../../../torchvision && uv pip install -e . --no-build-isolation)
(cd ../../../torchdata && uv pip install -e .)
(cd ../../../torchaudio && uv pip install -e . --no-build-isolation)
(cd ../../../FBGEMM/fbgemm_gpu && uv pip install -r requirements.txt && uv pip install -e . --no-build-isolation)
(cd ../../../torchrec && uv pip install -e .)
(cd ../../../detectron2 && uv pip install -e . --no-build-isolation)
(cd ../../../torchbenchmark && python install.py --continue_on_fail)
(cd ../../../torchbenchmark && python install.py --continue_on_fail $(if $(TORCHBENCH_MODELS),models $(TORCHBENCH_MODELS)))
uv pip uninstall torchrec-nightly fbgemm-gpu-nightly