diff --git a/onnxruntime/python/tools/transformers/models/sam2/README.md b/onnxruntime/python/tools/transformers/models/sam2/README.md index 83c0c51f09..b0d35ac79f 100644 --- 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 index 74048f9042..94e57ecb89 100644 --- 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 index 8ad69dee0a..40c408e851 100644 --- 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.