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
This commit is contained in:
Tianlei Wu 2024-09-26 20:57:03 -07:00 committed by GitHub
parent 3846f84218
commit ff8a48ef3b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 28 additions and 7 deletions

View file

@ -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.

View file

@ -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"

View file

@ -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.