Fix SAM2 benchmark script on cuda graph (#22377)

### Description
Update segment anything 2 benchmark script:
(1) Fix cuda graph in benchmark. Make sure --use_cuda_graph takes effect
and random_inputs() generates according to the dtype of the model.
(2) Add a parameter to enable profiling.
(3) Use latest cuda 12.6.2 and cudnn 9.5.
(4) Update README.md.

### Motivation and Context
Previous, --use_cuda_graph does not take effect. This fixes the
benchmark.
This commit is contained in:
Tianlei Wu 2024-10-10 08:46:44 -07:00 committed by GitHub
parent 1bc546af61
commit 2584d02c5e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 169 additions and 219 deletions

View file

@ -89,13 +89,15 @@ It is able to run demo on optimized model as well. For example,
python3 convert_to_onnx.py --sam2_dir $sam2_dir --optimize --dtype fp16 --use_gpu --demo
```
## Benchmark
## Benchmark and Profiling
We can create a conda environment then run GPU benchmark like the following:
```bash
conda create -n sam2_gpu python=3.11 -y
conda activate sam2_gpu
bash benchmark_sam2.sh $HOME gpu
install_dir=$HOME
profiling=true
bash benchmark_sam2.sh $install_dir gpu $profiling
```
or create a new conda environment for CPU benchmark:
@ -107,13 +109,14 @@ bash benchmark_sam2.sh $HOME cpu
The first parameter is a directory to clone git repositories or install CUDA/cuDNN for benchmark.
The second parameter can be either "gpu" or "cpu", which indicates the device to run benchmark.
The third parameter is optional. Value "true" will enable profiling after running benchmarking on GPU.
The script will automatically install required packages in current conda environment, download checkpoints, export onnx,
and run demo, benchmark and profiling.
and run demo, benchmark and optionally run profiling.
* The performance test result is in sam2_gpu.csv or sam2_cpu.csv, which can be loaded into Excel.
* The demo output is sam2_demo_fp16_gpu.png or sam2_demo_fp32_cpu.png.
* The profiling results are in *.nsys-rep or *.json files in current directory.
* The profiling results are in *.nsys-rep or *.json files in current directory. Use Nvidia NSight System to view the *.nsys-rep file.
## Limitations
- The exported image_decoder model does not support batch mode for now.

View file

@ -90,26 +90,23 @@ class TestConfig:
else:
return decoder_shape_dict(self.height, self.width, self.num_labels, self.num_points, self.num_masks)
def random_inputs(self):
def random_inputs(self) -> Mapping[str, torch.Tensor]:
dtype = self.dtype
if self.component == "image_encoder":
return {
"image": torch.randn(
self.batch_size, 3, self.height, self.width, dtype=torch.float32, device=self.device
)
}
return {"image": torch.randn(self.batch_size, 3, self.height, self.width, dtype=dtype, device=self.device)}
else:
return {
"image_features_0": torch.rand(1, 32, 256, 256, dtype=torch.float32, device=self.device),
"image_features_1": torch.rand(1, 64, 128, 128, dtype=torch.float32, device=self.device),
"image_embeddings": torch.rand(1, 256, 64, 64, dtype=torch.float32, device=self.device),
"image_features_0": torch.rand(1, 32, 256, 256, dtype=dtype, device=self.device),
"image_features_1": torch.rand(1, 64, 128, 128, dtype=dtype, device=self.device),
"image_embeddings": torch.rand(1, 256, 64, 64, dtype=dtype, device=self.device),
"point_coords": torch.randint(
0, 1024, (self.num_labels, self.num_points, 2), dtype=torch.float32, device=self.device
0, 1024, (self.num_labels, self.num_points, 2), dtype=dtype, device=self.device
),
"point_labels": torch.randint(
0, 1, (self.num_labels, self.num_points), dtype=torch.int32, device=self.device
),
"input_masks": torch.zeros(self.num_labels, 1, 256, 256, dtype=torch.float32, device=self.device),
"has_input_masks": torch.ones(self.num_labels, dtype=torch.float32, device=self.device),
"input_masks": torch.zeros(self.num_labels, 1, 256, 256, dtype=dtype, device=self.device),
"has_input_masks": torch.ones(self.num_labels, dtype=dtype, device=self.device),
"original_image_size": torch.tensor([self.height, self.width], dtype=torch.int32, device=self.device),
}
@ -314,7 +311,7 @@ def run_test(
width=args.width,
device=device,
use_tf32=True,
enable_cuda_graph=False,
enable_cuda_graph=enable_cuda_graph,
dtype=dtypes[args.dtype],
prefer_nhwc=args.prefer_nhwc,
repeats=args.repeats,

View file

@ -1,48 +1,35 @@
#!/bin/bash
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
# -------------------------------------------------------------------------
# Here assumes that we are using conda (Anaconda/Miniconda/Miniforge) environment.
# For example, you can create a new conda environment like the following before running this script:
# conda create -n sam2_gpu python=3.11 -y
# conda activate sam2_gpu
# bash benchmark_sam2.sh $HOME gpu
# Or create a new conda environment for CPU benchmark:
# conda create -n sam2_cpu python=3.11 -y
# conda activate sam2_cpu
# bash benchmark_sam2.sh $HOME cpu
# Please refer to README.md for the prerequisites and usage of this script.
# bash benchmark_sam2.sh <install_dir> <cpu|gpu> [profiling]
python=$CONDA_PREFIX/bin/python3
python="$CONDA_PREFIX/bin/python3"
# Directory of the script
dir="$( cd "$( dirname "$0" )" && pwd )"
# Directory of the script and ONNX models
dir="$(cd "$(dirname "$0")" && pwd)"
onnx_dir="$dir/sam2_onnx_models"
# Directory of the onnx models
onnx_dir=$dir/sam2_onnx_models
# Installation directory (default: $HOME)
install_dir="${1:-$HOME}"
# Directory to install CUDA, cuDNN, and git clone sam2 or onnxruntime source code.
install_dir=$HOME
if [ $# -ge 1 ]; then
install_dir=$1
fi
if ! [ -d $install_dir ]; then
echo "install_dir: $install_dir does not exist."
if [ ! -d "$install_dir" ]; then
echo "Error: install_dir '$install_dir' does not exist."
exit 1
fi
# Directory of the sam2 code by "git clone https://github.com/facebookresearch/segment-anything-2"
sam2_dir=$install_dir/segment-anything-2
# SAM2 code directory and model to benchmark
sam2_dir="$install_dir/segment-anything-2"
model="sam2_hiera_large"
# model name to benchmark
model=sam2_hiera_large
# Default to use GPU if available.
cpu_or_gpu="gpu"
if [ $# -ge 2 ] && ([ "$2" = "gpu" ] || [ "$2" = "cpu" ]); then
cpu_or_gpu=$2
# Default to GPU, switch to CPU if specified
cpu_or_gpu="${2:-gpu}"
if [ "$cpu_or_gpu" != "gpu" ] && [ "$cpu_or_gpu" != "cpu" ]; then
echo "Invalid option: $2. Please specify 'cpu' or 'gpu'."
exit 1
fi
echo "install_dir: $install_dir"
@ -51,148 +38,156 @@ echo "cpu_or_gpu: $cpu_or_gpu"
install_cuda_12()
{
pushd $install_dir
wget https://developer.download.nvidia.com/compute/cuda/12.5.1/local_installers/cuda_12.5.1_555.42.06_linux.run
sh cuda_12.5.1_555.42.06_linux.run --toolkit --toolkitpath=$install_dir/cuda12.5 --silent --override --no-man-page
wget https://developer.download.nvidia.com/compute/cuda/12.6.2/local_installers/cuda_12.6.2_560.35.03_linux.run
sh cuda_12.6.2_560.35.03_linux.run --toolkit --toolkitpath=$install_dir/cuda12.6 --silent --override --no-man-page
export PATH="$install_dir/cuda12.5/bin:$PATH"
export LD_LIBRARY_PATH="$install_dir/cuda12.5/lib64:$LD_LIBRARY_PATH"
export PATH="$install_dir/cuda12.6/bin:$PATH"
export LD_LIBRARY_PATH="$install_dir/cuda12.6/lib64:$LD_LIBRARY_PATH"
popd
}
install_cudnn_9()
{
pushd $install_dir
wget https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-9.4.0.58_cuda12-archive.tar.xz
mkdir $install_dir/cudnn9.4
tar -Jxvf cudnn-linux-x86_64-9.4.0.58_cuda12-archive.tar.xz -C $install_dir/cudnn9.4 --strip=1 --no-overwrite-dir
export LD_LIBRARY_PATH="$install_dir/cudnn9.4/lib:$LD_LIBRARY_PATH"
# Function to install cuDNN 9.4
install_cudnn_9() {
pushd "$install_dir"
wget -q https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-9.5.0.50_cuda12-archive.tar.xz
mkdir -p "$install_dir/cudnn9.5"
tar -Jxvf cudnn-linux-x86_64-9.5.0.50_cuda12-archive.tar.xz -C "$install_dir/cudnn9.5" --strip=1
export LD_LIBRARY_PATH="$install_dir/cudnn9.5/lib:$LD_LIBRARY_PATH"
popd
}
install_gpu()
{
if ! [ -d $install_dir/cuda12.5 ]; then
install_cuda_12
fi
if ! [ -d $install_dir/cudnn9.4 ]; then
install_cudnn_9
fi
# Install GPU dependencies
install_gpu() {
[ ! -d "$install_dir/cuda12.6" ] && install_cuda_12
[ ! -d "$install_dir/cudnn9.5" ] && install_cudnn_9
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
pip install onnxruntime-gpu onnx opencv-python matplotlib
}
install_cpu()
{
# Install CPU dependencies
install_cpu() {
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
pip install onnxruntime onnx opencv-python matplotlib
}
install_sam2()
{
pushd $install_dir
if ! [ -d $install_dir/segment-anything-2 ]; then
# Clone and install SAM2 if not already installed
install_sam2() {
pushd "$install_dir"
if [ ! -d "$sam2_dir" ]; then
git clone https://github.com/facebookresearch/segment-anything-2.git
fi
cd segment-anything-2
if pip show SAM-2 > /dev/null 2>&1; then
echo "SAM-2 is already installed."
else
pip install -e .
fi
if ! [ -f checkpoints/sam2_hiera_large.pt ]; then
echo "Downloading checkpoints..."
cd checkpoints
sh ./download_ckpts.sh
fi
cd "$sam2_dir"
pip show SAM-2 > /dev/null 2>&1 || pip install -e .
[ ! -f checkpoints/sam2_hiera_large.pt ] && (cd checkpoints && sh ./download_ckpts.sh)
popd
}
download_test_image()
{
if ! [ -f truck.jpg ]; then
curl https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/notebooks/images/truck.jpg > truck.jpg
# Download test image if not available
download_test_image() {
[ ! -f truck.jpg ] && curl -sO https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/notebooks/images/truck.jpg
}
run_cpu_benchmark() {
local repeats="$1"
$python convert_to_onnx.py --sam2_dir "$sam2_dir" --optimize --demo
for component in image_encoder image_decoder; do
$python benchmark_sam2.py --model_type "$model" --engine torch --sam2_dir "$sam2_dir" --repeats "$repeats" --dtype fp32 --component "$component"
# Run ONNX Runtime on exported model (not optimized)
$python benchmark_sam2.py --model_type "$model" --engine ort --sam2_dir "$sam2_dir" --repeats "$repeats" --onnx_path "${onnx_dir}/${model}_${component}.onnx" --dtype fp32 --component "$component"
# Run ONNX Runtime on optimized model
$python benchmark_sam2.py --model_type "$model" --engine ort --sam2_dir "$sam2_dir" --repeats "$repeats" --onnx_path "${onnx_dir}/${model}_${component}_fp32_cpu.onnx" --dtype fp32 --component "$component"
done
}
run_gpu_benchmark() {
local repeats="$1"
$python convert_to_onnx.py --sam2_dir "$sam2_dir" --optimize --use_gpu --dtype fp32
$python convert_to_onnx.py --sam2_dir "$sam2_dir" --optimize --use_gpu --dtype fp16 --demo
for component in image_encoder image_decoder; do
for dtype in bf16 fp32 fp16; do
$python benchmark_sam2.py --model_type "$model" --engine torch --sam2_dir "$sam2_dir" --repeats "$repeats" --use_gpu --dtype $dtype --component "$component"
done
done
component="image_encoder"
for dtype in fp32 fp16; do
#TODO: --prefer_nhwc does not help with performance
$python benchmark_sam2.py --model_type "$model" --engine ort --sam2_dir "$sam2_dir" --repeats "$repeats" --use_gpu --dtype $dtype --component "$component" --onnx_path "${onnx_dir}/${model}_${component}_${dtype}_gpu.onnx" --use_cuda_graph
done
component="image_decoder"
for dtype in fp32 fp16; do
# TODO: decoder does not work with cuda graph
$python benchmark_sam2.py --model_type "$model" --engine ort --sam2_dir "$sam2_dir" --repeats "$repeats" --use_gpu --dtype $dtype --component "$component" --onnx_path "${onnx_dir}/${model}_${component}_${dtype}_gpu.onnx"
done
}
run_torch_compile_gpu_benchmark() {
local repeats="$1"
# Test different torch compile modes on image encoder
for torch_compile_mode in none max-autotune reduce-overhead max-autotune-no-cudagraphs
do
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir "$sam2_dir" --repeats "$repeats" --use_gpu --dtype fp16 --component image_encoder --torch_compile_mode $torch_compile_mode
done
}
# Main script
run_benchmarks() {
if [ ! -v CONDA_PREFIX ]; then
echo "Please activate conda environment before running this script."
exit 1
fi
# Install dependencies
[ "$cpu_or_gpu" = "gpu" ] && install_gpu || install_cpu
install_sam2
download_test_image
# Run benchmarks
output_csv="sam2_${cpu_or_gpu}.csv"
if [ ! -f "$output_csv" ]; then
echo "Running $cpu_or_gpu benchmark..."
if [ "$cpu_or_gpu" = "gpu" ]; then
run_gpu_benchmark 1000
run_torch_compile_gpu_benchmark 1000
else
run_cpu_benchmark 100
fi
cat benchmark*.csv > combined_csv
awk '!x[$0]++' combined_csv > "$output_csv"
rm combined_csv
echo "Benchmark results saved in $output_csv"
else
echo "$output_csv already exists, skipping benchmark..."
fi
}
run_cpu()
{
repeats=$1
run_benchmarks
$python convert_to_onnx.py --sam2_dir $sam2_dir --optimize --demo
#--------------------------------------------------------------------------
# Below are for profiling
#--------------------------------------------------------------------------
echo "Benchmarking SAM2 model $model image encoder for PyTorch ..."
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir $sam2_dir --repeats $repeats --dtype fp32
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir $sam2_dir --repeats $repeats --dtype fp16
echo "Benchmarking SAM2 model $model image encoder for PyTorch ..."
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir $sam2_dir --repeats $repeats --dtype fp32 --component image_decoder
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir $sam2_dir --repeats $repeats --dtype fp16 --component image_decoder
echo "Benchmarking SAM2 model $model image encoder for ORT ..."
$python benchmark_sam2.py --model_type $model --engine ort --sam2_dir $sam2_dir --repeats $repeats --onnx_path ${onnx_dir}/${model}_image_encoder.onnx --dtype fp32
$python benchmark_sam2.py --model_type $model --engine ort --sam2_dir $sam2_dir --repeats $repeats --onnx_path ${onnx_dir}/${model}_image_encoder_fp32_cpu.onnx --dtype fp32
echo "Benchmarking SAM2 model $model image decoder for ORT ..."
$python benchmark_sam2.py --model_type $model --engine ort --sam2_dir $sam2_dir --repeats $repeats --onnx_path ${onnx_dir}/${model}_image_decoder.onnx --component image_decoder
$python benchmark_sam2.py --model_type $model --engine ort --sam2_dir $sam2_dir --repeats $repeats --onnx_path ${onnx_dir}/${model}_image_decoder_fp32_cpu.onnx --component image_decoder
}
run_gpu()
{
repeats=$1
$python convert_to_onnx.py --sam2_dir $sam2_dir --optimize --use_gpu --dtype fp32
$python convert_to_onnx.py --sam2_dir $sam2_dir --optimize --use_gpu --dtype fp16 --demo
echo "Benchmarking SAM2 model $model image encoder for PyTorch ..."
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir $sam2_dir --repeats $repeats --use_gpu --dtype bf16
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir $sam2_dir --repeats $repeats --use_gpu --dtype fp32
# Test different torch compile modes on image encoder (none will disable compile and use eager mode).
for torch_compile_mode in none max-autotune reduce-overhead max-autotune-no-cudagraphs
do
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir $sam2_dir --repeats $repeats --use_gpu --dtype fp16 --component image_encoder --torch_compile_mode $torch_compile_mode
done
echo "Benchmarking SAM2 model $model image decoder for PyTorch ..."
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir $sam2_dir --repeats $repeats --use_gpu --dtype bf16 --component image_decoder
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir $sam2_dir --repeats $repeats --use_gpu --dtype fp32 --component image_decoder
$python benchmark_sam2.py --model_type $model --engine torch --sam2_dir $sam2_dir --repeats $repeats --use_gpu --dtype fp16 --component image_decoder
echo "Benchmarking SAM2 model $model image encoder for ORT ..."
$python benchmark_sam2.py --model_type $model --engine ort --sam2_dir $sam2_dir --repeats $repeats --onnx_path ${onnx_dir}/${model}_image_encoder_fp16_gpu.onnx --use_gpu --dtype fp16
$python benchmark_sam2.py --model_type $model --engine ort --sam2_dir $sam2_dir --repeats $repeats --onnx_path ${onnx_dir}/${model}_image_encoder_fp32_gpu.onnx --use_gpu --dtype fp32
echo "Benchmarking SAM2 model $model image decoder for ORT ..."
$python benchmark_sam2.py --model_type $model --engine ort --sam2_dir $sam2_dir --repeats $repeats --onnx_path ${onnx_dir}/${model}_image_decoder_fp16_gpu.onnx --component image_decoder --use_gpu --dtype fp16
$python benchmark_sam2.py --model_type $model --engine ort --sam2_dir $sam2_dir --repeats $repeats --onnx_path ${onnx_dir}/${model}_image_decoder_fp32_gpu.onnx --component image_decoder --use_gpu
}
# Build onnxruntime-gpu from source for profiling.
build_onnxruntime_gpu_for_profiling()
{
pushd $install_dir
# Build onnxruntime-gpu from source for profiling
build_onnxruntime_gpu_for_profiling() {
pushd "$install_dir"
if ! [ -d onnxruntime ]; then
git clone https://github.com/microsoft/onnxruntime
fi
cd onnxruntime
# Get the CUDA compute capability of the GPU.
CUDA_ARCH=$(python3 -c "import torch; cc = torch.cuda.get_device_capability(); print(f'{cc[0]}{cc[1]}')")
if [ -n "$CUDA_ARCH" ]; then
pip install --upgrade pip cmake psutil setuptools wheel packaging ninja numpy==1.26.4
sh build.sh --config Release --build_dir build/cuda12 --build_shared_lib --parallel \
--use_cuda --cuda_version 12.5 --cuda_home $install_dir/cuda12.5 \
--cudnn_home $install_dir/cudnn9.4 \
--use_cuda --cuda_version 12.6 --cuda_home $install_dir/cuda12.6 \
--cudnn_home $install_dir/cudnn9.5 \
--build_wheel --skip_tests \
--cmake_generator Ninja \
--compile_no_warning_as_error \
@ -203,33 +198,24 @@ build_onnxruntime_gpu_for_profiling()
pip install build/cuda12/Release/dist/onnxruntime_gpu-*-linux_x86_64.whl numpy==1.26.4
else
echo "PyTorch is not installed or No CUDA device found."
echo "No CUDA device found."
exit 1
fi
popd
}
# Run profiling with NVTX.
run_nvtx_profile()
{
pip install nvtx cuda-python==12.5.0
pip install nvtx cuda-python==12.6.0
# Only trace one device to avoid huge output file size.
device_id=0
# Environment variables
envs="CUDA_VISIBLE_DEVICES=$device_id,ORT_ENABLE_CUDNN_FLASH_ATTENTION=1,LD_LIBRARY_PATH=$LD_LIBRARY_PATH"
# For cuda graphs, node activities will be collected and CUDA graphs will not be traced as a whole.
# This may cause significant runtime overhead. But it is useful to understand the performance of individual nodes.
cuda_graph_trace=node
for engine in ort torch
do
for component in image_encoder image_decoder
do
sudo $install_dir/cuda12.5/bin/nsys profile --capture-range=nvtx --nvtx-capture='one_run' \
for engine in ort torch; do
for component in image_encoder image_decoder; do
sudo $install_dir/cuda12.6/bin/nsys profile --capture-range=nvtx --nvtx-capture='one_run' \
--gpu-metrics-device $device_id --force-overwrite true \
--sample process-tree --backtrace fp --stats true \
-t cuda,cudnn,cublas,osrt,nvtx --cuda-memory-usage true --cudabacktrace all \
@ -246,10 +232,8 @@ run_nvtx_profile()
}
# Run profiling with PyTorch
run_torch_profile()
{
for component in image_encoder image_decoder
do
run_torch_profile() {
for component in image_encoder image_decoder; do
$python benchmark_sam2.py --model_type $model --engine torch \
--sam2_dir $sam2_dir --warm_up 1 --repeats 0 \
--component $component \
@ -257,50 +241,16 @@ run_torch_profile()
done
}
if ! [ -v CONDA_PREFIX ]; then
echo "Please activate conda environment before running this script."
exit 1
fi
run_profilings() {
build_onnxruntime_gpu_for_profiling
# Check whether nvidia-smi is available to determine whether to install GPU or CPU version.
if [ "$cpu_or_gpu" = "gpu" ]; then
install_gpu
else
install_cpu
fi
install_sam2
download_test_image
if ! [ -f sam2_${cpu_or_gpu}.csv ]; then
if [ "$cpu_or_gpu" = "gpu" ]; then
echo "Running GPU benchmark..."
run_gpu 1000
else
echo "Running CPU benchmark..."
run_cpu 100
fi
cat benchmark*.csv > combined_csv
awk '!x[$0]++' combined_csv > sam2_${cpu_or_gpu}.csv
rm combined_csv
echo "Benchmarking SAM2 model $model results are saved in sam2_${cpu_or_gpu}.csv"
else
echo "sam2_${cpu_or_gpu}.csv already exists, skipping benchmarking..."
fi
if [ "$cpu_or_gpu" = "gpu" ]; then
echo "Running GPU profiling..."
if ! [ -f sam2_fp16_profile_image_decoder_ort_${cpu_or_gpu}.nsys-rep ]; then
rm -f *.nsys-rep
rm -f *.sqlite
build_onnxruntime_gpu_for_profiling
run_nvtx_profile
else
echo "sam2_fp16_profile_image_decoder_ort_${cpu_or_gpu}.nsys-rep already exists, skipping GPU profiling..."
fi
rm -f *.nsys-rep *.sqlite
run_nvtx_profile
run_torch_profile
}
profiling="${3:-false}"
if [ "$profiling" = "true" ] && [ "$cpu_or_gpu" = "gpu" ]; then
run_profilings
fi