EdgeSAM
用于SAM设备部署的闭环提示蒸馏
(*通讯作者)
[论文
]
[项目主页
]
[Hugging Face演示
]
[iOS应用
]
https://github.com/chongzhou96/EdgeSAM/assets/15973859/fe1cd104-88dc-4690-a5ea-ff48ae013db3
更新
- 2024/07/23: 我们发布了训练和评估代码,请查看 README_TRAIN.md。
- 2024/06/05: 查看我们由EdgeSAM驱动的iOS应用CutCha。
- 2024/01/01: EdgeSAM已集成到X-AnyLabeling中。
- 2023/12/19: EdgeSAM现已在ISAT中得到支持,这是一个分割标注工具。
- 2023/12/16: EdgeSAM现已在Grounded-Segment-Anything中得到支持。查看grounded-edge-sam演示。感谢IDEA Research团队!
- 2023/12/14: autodistill-grounded-edgesam结合了Grounding DINO和EdgeSAM创建了Grounded EdgeSAM [博客]。感谢Roboflow团队!
- 2023/12/13: 添加ONNX导出功能,并使用ONNX作为后端加速网页演示。
概述
EdgeSAM是分割任意物体模型(SAM)的加速变体,针对边缘设备上的高效执行进行了优化,性能损失最小。 与原始SAM相比,它实现了40倍的速度提升,并且优于MobileSAM,在边缘设备上部署时快14倍,同时在COCO和LVIS上的mIoU分别提高了2.3和3.2。 EdgeSAM还是第一个能在iPhone 14上以超过30 FPS运行的SAM变体。
在这张图中,我们展示了EdgeSAM与SAM和MobileSAM相比的编码器吞吐量,以及在SA-1K数据集(从SA-1B采样)上使用框和点提示的mIoU性能。
方法
我们的方法涉及将原始基于ViT的SAM图像编码器蒸馏成纯CNN架构,更适合边缘设备。我们仔细评估了各种蒸馏策略,并证明任务无关的编码器蒸馏无法捕获SAM中包含的全部知识。为了克服这一瓶颈,我们在蒸馏过程中包含了提示编码器和掩码解码器,并在闭环中使用框和点提示,使得蒸馏模型能够准确捕捉用户输入和掩码生成之间的复杂动态。
性能
方法 | 训练集 | COCO AP | COCO APs | COCO APm | COCO APl | GFLops | MParam. | FPS iPhone 14 | FPS 2080 Ti | FPS 3090 |
---|---|---|---|---|---|---|---|---|---|---|
SAM | SA-1B | 46.1 | 33.6 | 51.9 | 57.7 | 2734.8 | 641.1 | - | 4.3 | - |
FastSAM | 2% SA-1B | 37.9 | 23.9 | 43.4 | 50.0 | 887.6 | 68.2 | - | - | 25.0* |
MobileSAM | 1% SA-1B | 39.4 | 26.9 | 44.4 | 52.2 | 38.2 | 9.8 | 4.9 | 103.5 | 100.0* |
EdgeSAM | 1% SA-1B | 42.2 | 29.6 | 47.6 | 53.9 | 22.1 | 9.6 | 38.7 | 164.3 | - |
EdgeSAM-3x | 3% SA-1B | 42.7 | 30.0 | 48.6 | 54.5 | 22.1 | 9.6 | 38.7 | 164.3 | - |
EdgeSAM-10x | 10% SA-1B | 43.0 | 30.3 | 48.9 | 55.1 | 22.1 | 9.6 | 38.7 | 164.3 | - |
在此表中,我们报告了COCO数据集上的掩码mAP。使用ViTDet-H作为检测器(其框mAP为58.7)提供框提示。对于速度基准测试,我们同时推断编码器和解码器(使用单个提示)。FLOPs基于1024x1024输入分辨率计算。带的数字复制自MobileSAM。3x和10x表示使用更多数据进行训练。这里,我们没有应用额外的掩码细化迭代,遵循原始SAM论文的设置。*
目录
安装
代码需要python>=3.8
,我们使用torch==2.0.0
和torchvision==0.15.1
。请参考
官方PyTorch安装说明。
- 在本地克隆仓库:
git clone https://github.com/chongzhou96/EdgeSAM.git && cd EdgeSAM
- 安装额外的依赖项:
pip install -r requirements.txt
- 安装EdgeSAM:
pip install -e .
使用
- 下载检查点(请参考检查点了解更多关于PyTorch和CoreML检查点的详细信息):
mkdir weights
wget -P weights/ https://huggingface.co/spaces/chongzhou/EdgeSAM/resolve/main/weights/edge_sam.pth
wget -P weights/ https://huggingface.co/spaces/chongzhou/EdgeSAM/resolve/main/weights/edge_sam_3x.pth
- 您可以轻松地将EdgeSAM集成到您的Python代码中,只需几行代码:
from edge_sam import SamPredictor, sam_model_registry
sam = sam_model_registry["edge_sam"](checkpoint="<path/to/checkpoint>")
predictor = SamPredictor(sam)
predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(<input_prompts>)
由于EdgeSAM遵循与SAM相同的编码器-解码器架构,它们的用法非常相似。一个小的区别是EdgeSAM允许为每个提示输出1、3和4个掩码候选,而SAM只能输出1或3个掩码。有关更多详细信息,请参考示例Jupyter Notebook。
训练和评估
请参考README_TRAIN.md了解更多详细信息。
网页演示
安装EdgeSAM并下载检查点后,您可以使用以下命令启动交互式网页演示:
python web_demo/gradio_app.py
默认情况下,演示托管在http://0.0.0.0:8080/
上,并期望edge_sam_3x.pth
存储在weights/
文件夹中。您可以通过以下方式更改默认行为:
python web_demo/gradio_app.py --checkpoint [CHECKPOINT] --server-name [SERVER_NAME] --port [PORT]
由于EdgeSAM可以在手机上流畅运行,即使您没有GPU也没关系。
我们已经在Hugging Face Space上部署了相同的网页演示[链接]。然而,由于它使用CPU作为后端并由所有用户共享,体验可能不如本地部署那么好。 非常感谢Hugging Face团队为我们提供GPU支持!
使用ONNX后端加速网页演示
-
如果您的机器没有GPU,请安装onnxruntime,使用
pip install onnxruntime
;如果有GPU,则使用pip install onnxruntime-gpu
(但不要同时安装两者)。我们的实现在版本1.16.3
下进行了测试。 -
将ONNX模型下载到
weights/
文件夹:
wget -P weights/ https://huggingface.co/spaces/chongzhou/EdgeSAM/resolve/main/weights/edge_sam_3x_encoder.onnx
wget -P weights/ https://huggingface.co/spaces/chongzhou/EdgeSAM/resolve/main/weights/edge_sam_3x_decoder.onnx
- 启动演示:
python web_demo/gradio_app.py --enable-onnx
- 在浏览器中访问 http://0.0.0.0:8080。
CoreML / ONNX 导出
CoreML
我们提供了一个脚本,可以将训练好的EdgeSAM PyTorch模型导出为两个CoreML模型包,一个用于编码器,另一个用于解码器。您也可以在检查点处下载导出的CoreML模型。
对于编码器:
python scripts/export_coreml_model.py [CHECKPOINT]
对于解码器:
python scripts/export_coreml_model.py [CHECKPOINT] --decoder --use-stability-score
由于EdgeSAM没有对原始SAM的IoU令牌执行知识蒸馏,其IoU预测可能不太可靠。因此,我们使用稳定性分数来代替选择掩码。如果您想坚持使用IoU预测,可以删除--use-stability-score
。
以下显示了在iPhone 14上由Xcode测量的EdgeSAM CoreML模型的性能报告(左:编码器,右:解码器):
已知问题和模型描述
在coremltools==7.1
版本中,您可能在导出过程中遇到断言错误,例如assert len(inputs) <= 3 or inputs[3] is None
。一种解决方法是按照回溯路径注释掉这个断言,例如/opt/anaconda3/envs/EdgeSAM/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/ops.py 第1573行
。
由于CoreML不支持动态目标大小的插值,转换后的CoreML模型不包含预处理(即调整大小-归一化-填充)和后处理(即调整回原始大小)。
编码器接受1x3x1024x1024
的图像作为输入,输出1x256x64x64
的图像嵌入。然后解码器将图像嵌入与点坐标和点标签作为输入。点坐标遵循(height, width)
格式,左上角为(0, 0)
。点标签的选项为0:负点
,1:正点
,2:框的左上角
,和3:框的右下角
。
ONNX
类似于CoreML导出,您可以使用以下命令分别将编码器和解码器导出为ONNX模型:
对于编码器:
python scripts/export_onnx_model.py [CHECKPOINT]
对于解码器:
python scripts/export_onnx_model.py [CHECKPOINT] --decoder --use-stability-score
检查点
请从其Hugging Face Space下载EdgeSAM的检查点(所有EdgeSAM变体仅在训练图像数量上有所不同):
模型 | COCO mAP | PyTorch | CoreML | ONNX |
---|---|---|---|---|
SAM | 46.1 | - | - | - |
EdgeSAM | 42.1 | 下载 | [编码器] [解码器] | [编码器] [解码器] |
EdgeSAM-3x | 42.7 | 下载 | [编码器] [解码器] | [编码器] [解码器] |
EdgeSAM-10x | 43 | 待定 | 待定 | 待定 |
注意:使用前需要解压CoreML模型包。
iOS应用
我们计划将在现场演示中使用的iOS应用发布到App Store。敬请期待!
致谢
本研究得到了RIE2020产业对齐基金产业合作项目(IAF-ICP)资助计划的支持,以及行业合作伙伴的现金和实物贡献。我们感谢Han Soong Chong在演示应用方面所做的努力。
我们感谢以下项目,它们使EdgeSAM成为可能:SAM, MobileSAM, FastSAM, TinyViT, 和 RepViT。
引用
@article{zhou2023edgesam,
title={EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM},
author={Zhou, Chong and Li, Xiangtai and Loy, Chen Change and Dai, Bo},
journal={arXiv preprint arXiv:2312.06660},
year={2023}
}
许可证
本项目根据NTU S-Lab许可证1.0授权。重新分发和使用应遵循此许可证。