EdgeSAM

基于提示循环的蒸馏技术，用于在端侧设备上部署 SAM

周冲¹, 李翔泰¹, 陈昌毅·洛伊^1*, 戴博²

(*通讯作者)

¹S-Lab，南洋理工大学， ²上海人工智能实验室

[论文] [项目页面] [Hugging Face 演示] [iOS 应用]

https://github.com/chongzhou96/EdgeSAM/assets/15973859/fe1cd104-88dc-4690-a5ea-ff48ae013db3

观看完整直播演示视频：[YouTube] [Bilibili]

更新

• 2024年7月23日：我们发布了训练和评估代码，请查看 README_TRAIN.md。
• 2024年6月5日：请体验由 EdgeSAM 提供支持的 iOS 应用 CutCha。
• 2024年1月1日：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 是 Segment Anything Model (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 论文中的设置应用额外的掩码细化迭代。

目录

• 安装
• 使用
• 训练与评估
• Web 演示
• CoreML / ONNX 导出
• 检查点
• iOS 应用
• 致谢
• 引用
• 许可证

安装

代码需要 python>=3.8，我们使用 torch==2.0.0 和 torchvision==0.15.1。请参考 官方 PyTorch 安装说明。

1. 在本地克隆仓库：

git clone https://github.com/chongzhou96/EdgeSAM.git && cd EdgeSAM

2. 安装其他依赖项：

pip install -r requirements.txt

3. 安装 EdgeSAM：

pip install -e .

使用方法

1. 下载检查点文件（有关 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

2. 您可以通过以下几行代码轻松地将 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。

Web 演示

在安装 EdgeSAM 并下载检查点文件后，您可以使用以下命令启动交互式 Web 演示：

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 中部署了相同的 Web 演示 [链接]。 然而，由于它使用 CPU 作为后端且所有用户共享资源，体验可能不如本地部署。 非常感谢 Hugging Face 团队为我们提供 GPU 支持！

使用 ONNX 后端加速 Web 演示

1. 如果您的机器没有 GPU，请使用 pip install onnxruntime 安装 onnxruntime；如果有 GPU，则使用 pip install onnxruntime-gpu（但不要同时安装两者）。我们的实现已在版本 1.16.3 下测试通过。

2. 将 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

3. 启动演示：

python web_demo/gradio_app.py --enable-onnx

4. 在浏览器中访问 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 即可。

以下是使用 Xcode 在 iPhone 14 上测量的 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

检查点

请从 EdgeSAM 的 Hugging Face Space 下载检查点文件（所有 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={周冲、李向泰、陈昌礼、戴博},
  journal={arXiv 预印本 arXiv:2312.06660},
  year={2023}
}

许可证

本项目采用 NTU S-Lab 许可证 1.0 进行授权。再分发和使用应遵守该许可证的规定。

EdgeSAM 快速上手指南

EdgeSAM 是 Segment Anything Model (SAM) 的加速版本，专为边缘设备（如手机、嵌入式设备）优化。相比原版 SAM 速度提升 40 倍，在 iPhone 14 上可实现超过 30 FPS 的实时分割性能，同时保持了较高的分割精度。

环境准备

• 操作系统: Linux, macOS, Windows
• Python 版本: >= 3.8
• 核心依赖:
  • torch == 2.0.0
  • torchvision == 0.15.1
• 硬件建议: 支持 CPU 推理（适合移动端部署），若有 GPU 可加速训练和 Web Demo 体验。

注意: 请优先参考 PyTorch 官方安装指南 安装对应版本的 torch 和 torchvision。国内用户可使用清华源或阿里源加速安装：

pip install torch==2.0.0 torchvision==0.15.1 -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆仓库

   git clone https://github.com/chongzhou96/EdgeSAM.git && cd EdgeSAM
   

2. 安装依赖

   pip install -r requirements.txt
   

   (国内用户建议添加 -i https://pypi.tuna.tsinghua.edu.cn/simple)

3. 安装 EdgeSAM

   pip install -e .
   

4. 下载预训练权重 创建权重目录并下载模型文件（提供基础版和增强版）：

   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
   

   提示: 如果 wget 下载缓慢，可直接在浏览器访问上述链接下载，或将文件手动放入 weights/ 目录。推荐使用 edge_sam_3x.pth 以获得更好的精度。

基本使用

EdgeSAM 的使用方式与原始 SAM 高度兼容。以下是最简单的 Python 调用示例：

from edge_sam import SamPredictor, sam_model_registry

# 1. 加载模型 (替换为实际下载的权重路径)
sam = sam_model_registry["edge_sam"](checkpoint="weights/edge_sam_3x.pth")

# 2. 初始化预测器
predictor = SamPredictor(sam)

# 3. 设置输入图像 (支持 numpy array)
predictor.set_image(<your_image>)

# 4. 进行预测 (输入可以是点、框等 prompts)
# 示例：根据输入提示生成掩码
masks, scores, logits = predictor.predict(<input_prompts>)

关键特性说明：

• 输出数量: 与原版 SAM 不同，EdgeSAM 允许为每个提示输出 1、3 或 4 个候选掩码。
• 架构兼容: 采用相同的 Encoder-Decoder 架构，便于集成到现有 SAM 工作流中。

更多详细用法（如交互式标注、批量处理）请参考项目自带的 notebooks/predictor_example.ipynb。