Update SAM2 benchmark script and doc (#22238)

(1) Fix a bug of parameters order. (2) Update benchmark script: * download test image if not exist * combine multiple csv files into one file, and remove duplicated lines (3) Add a section for benchmark in README.md
microsoft · Sep 27, 2024 · ff8a48e · ff8a48e
1 parent 3846f84
commit ff8a48e
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 7 deletions.
diff --git a/onnxruntime/python/tools/transformers/models/sam2/README.md b/onnxruntime/python/tools/transformers/models/sam2/README.md
@@ -26,6 +26,7 @@ Clone the SAM2 git repository and download the checkpoints:
 ```bash
 git clone https://github.com/facebookresearch/segment-anything-2.git
 cd segment-anything-2
+export sam2_dir=$PWD
 python3 -m pip install -e .
 cd checkpoints
 sh ./download_ckpts.sh
@@ -42,7 +43,7 @@ curl https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.p
 ## Export ONNX
 To export ONNX models, run the convert_to_onnx.py script and specify the segment-anything-2 directory created by the above git clone command:
 ```bash
-python3 convert_to_onnx.py  --sam2_dir path/to/segment-anything-2
+python3 convert_to_onnx.py  --sam2_dir $sam2_dir
 ```
 
 The exported ONNX models will be found in the sam2_onnx_models sub-directory. You can change the output directory using the `--output_dir` option.
@@ -58,12 +59,12 @@ python3 convert_to_onnx.py  -h
 
 To optimize the onnx models for CPU with float32 data type:
 ```bash
-python3 convert_to_onnx.py  --sam2_dir path/to/segment-anything-2 --optimize --dtype fp32
+python3 convert_to_onnx.py  --sam2_dir $sam2_dir --optimize --dtype fp32
 ```
 
 To optimize the onnx models for GPU with float16 data type:
 ```bash
-python3 convert_to_onnx.py  --sam2_dir path/to/segment-anything-2 --optimize --dtype fp16 --use_gpu
+python3 convert_to_onnx.py  --sam2_dir $sam2_dir --optimize --dtype fp16 --use_gpu
 ```
 
 Another option is to use optimizer.py like the following:
@@ -80,13 +81,22 @@ The optimizer.py could be helpful when you have SAM2 onnx models that is exporte
 The exported ONNX models can run on a CPU. The demo will output sam2_demo.png.
 ```bash
 curl https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/notebooks/images/truck.jpg > truck.jpg
-python3 convert_to_onnx.py  --sam2_dir path/to/segment-anything-2 --demo
+python3 convert_to_onnx.py  --sam2_dir $sam2_dir --demo
 ```
 
 It is able to run demo on optimized model as well. For example,
 ```bash
-python3 convert_to_onnx.py  --sam2_dir path/to/segment-anything-2 --optimize --dtype fp16 --use_gpu --demo
+python3 convert_to_onnx.py  --sam2_dir $sam2_dir --optimize --dtype fp16 --use_gpu --demo
 ```
 
+## Benchmark
+To prepare an environment for benchmark, follow [Setup Environment](#setup-environment) and [Download Checkpoints](#download-checkpoints).
+
+Run the benchmark like the following:
+```bash
+sh benchmark_sam2.sh
+```
+The result is in sam2.csv, which can be loaded into Excel.
+
 ## Limitations
 - The exported image_decoder model does not support batch mode for now.
diff --git a/onnxruntime/python/tools/transformers/models/sam2/benchmark_sam2.sh b/onnxruntime/python/tools/transformers/models/sam2/benchmark_sam2.sh
@@ -11,7 +11,8 @@ dir="$( cd "$( dirname "$0" )" && pwd )"
 onnx_dir=$dir/sam2_onnx_models
 
 # Directory of the sam2 code by "git clone https://github.com/facebookresearch/segment-anything-2"
-sam2_dir=~/segment-anything-2
+# It reads from the sam2_dir environment variable, or defaults to ~/segment-anything-2.
+sam2_dir=${sam2_dir:-~/segment-anything-2}
 
 # model name to benchmark
 model=sam2_hiera_large
@@ -65,8 +66,18 @@ run_gpu()
     python3 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
 }
 
+if ! [ -f truck.jpg ]; then
+    curl https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/notebooks/images/truck.jpg > truck.jpg
+fi
+
 if python3 -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then
     run_gpu 1000
 else
     run_cpu 100
 fi
+
+cat benchmark*.csv > combined_csv
+awk '!x[$0]++' combined_csv > sam2.csv
+rm combined_csv
+
+echo "Benchmarking SAM2 model $model done. Results are saved in sam2.csv"
diff --git a/onnxruntime/python/tools/transformers/models/sam2/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/sam2/convert_to_onnx.py
@@ -125,7 +125,7 @@ def parse_arguments():
     return args
 
 
-def optimize_sam2_model(onnx_model_path, optimized_model_path, use_gpu: bool, float16: bool):
+def optimize_sam2_model(onnx_model_path, optimized_model_path, float16: bool, use_gpu: bool):
     print(f"Optimizing {onnx_model_path} to {optimized_model_path} with float16={float16} and use_gpu={use_gpu}...")
 
     # Import from source directory.