mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-06-26 03:00:54 +00:00
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:
parent
3846f84218
commit
ff8a48ef3b
3 changed files with 28 additions and 7 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Reference in a new issue