SAM(ONNX)-推理指导

概述

Segment Anything Model (SAM) 是由 Meta 开源的图像分割大模型，在计算机视觉领域（CV）取得了新的突破。SAM 可在不需要任何标注的情况下，对任何图像中的任何物体进行分割，SAM 的开源引起了业界的广泛反响，被称为计算机视觉领域的 GPT。

论文：

[Segment Anything](https://arxiv.org/abs/2304.02643)
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, Ross Girshick

参考实现：

url=https://github.com/facebookresearch/segment-anything.git
commit_id=6fdee8f2727f4506cfbbe553e23b895e27956588
model_name=sam_vit_b_01ec64

推理环境准备

该模型需要以下插件与驱动

表 1 版本配套表

配套	版本	环境准备指导
固件与驱动	25.2.RC1	Pytorch框架推理环境准备
CANN	8.2.RC1	-
Python	3.11.10	-
PyTorch	2.1.0	-
说明：仅支持Atlas 300I Duo 推理卡，请以CANN版本选择实际固件与驱动版本。	\	\

快速上手

1. 获取源码

git clone https://gitee.com/ascend/ModelZoo-PyTorch.git
cd ModelZoo-PyTorch/ACL_PyTorch/built-in/cv/SAM
git clone https://github.com/facebookresearch/segment-anything.git
cd segment-anything
git reset --hard 6fdee8f2727f4506cfbbe553e23b895e27956588
git apply ../segment_anything_diff.patch
pip3 install -e .
cd ..

2. 安装依赖。

安装基础环境。

pip3 install -r requirements.txt

说明：如果某些库通过此方式安装失败，可使用 pip3 install 单独进行安装。

安装 msit 的 surgeon 组件和 benchmark 组件。

3. 输入输出数据描述

SAM 首先会自动分割图像中的所有内容，但是如果你需要分割某一个目标物体，则需要你输入一个目标物体上的坐标，比如一张图片你想让SAM分割Cat或Dog这个目标的提示坐标，SAM会自动在照片中猫或狗进行分割，在离线推理时，会转成encoder模型和decoder模型，其输入输出详情如下：

encoder输入数据

输入数据数据类型大小数据排布格式
x FLOAT32 1 x 3 x 1024 x 1024 NCHW
encoder输出数据

输出数据数据类型大小数据排布格式
image_embeddings FLOAT32 1 x 256 x 64 x 64 NCHW

输入数据	数据类型	大小	数据排布格式
x	FLOAT32	1 x 3 x 1024 x 1024	NCHW

decoder输入数据

输入数据	数据类型	大小	数据排布格式
image_embeddings	FLOAT32	1 x 256 x 64 x 64	NCHW
point_coords	FLOAT32	1 x -1 x 2	ND
point_labels	FLOAT32	1 x -1	ND
mask_input	FLOAT32	1 x 1 x 256 x 256	NCHW
has_mask_input	FLOAT32	1	ND

decoder输出数据

输出数据数据类型大小数据排布格式
iou_predictions FLOAT32 -1 x 1 ND
low_res_masks FLOAT32 -1 x 1 x -1 x -1 ND

输出数据	数据类型	大小	数据排布格式
iou_predictions	FLOAT32	-1 x 1	ND
low_res_masks	FLOAT32	-1 x 1 x -1 x -1	ND

4. 准备数据集

取仓库里的 demo 图片进行端到端测试。

mkdir data
cd data
wget -O demo.jpg https://raw.githubusercontent.com/facebookresearch/segment-anything/571794162e0887c15d12b809505b902c7bf8b4db/notebooks/images/truck.jpg
cd ..

下载coco2017数据集进行精度测试。

下载并解压COCO-2017数据集的图片与标注，放置coco2017目录下

coco2017
├── annotations/
│   └── instances_val2017.json
└── val2017/
    ├── 000000000139.jpg
    ├── 000000000139.jpg
    └── ...

5. 模型转换

5.1 获取权重文件

GitHub 仓库提供了三种大小的权重文件：vit_h、vit_l、vit_b。这里以 vit_b 为例。

mkdir models
cd models
wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
cd ..

5.2 导出 ONNX 模型

python3 segment-anything/scripts/export_onnx_model.py  \
  --checkpoint models/sam_vit_b_01ec64.pth  \
  --model-type vit_b  \
  --opset 14 \
  --encoder-output models/encoder.onnx \
  --decoder-output models/decoder.onnx \
  --return-single-mask

参数说明：

checkpoint：模型权重文件路径。
model-type：模型类型。
opset：onnx算子集版本。
encoder-output：保存encoder模型的输出ONNX模型的文件路径。
decoder-output：保存decoder模型的输出ONNX模型的文件路径。
return-single-mask：设置最优mask模式。

5.3 使用 onnxsim 简化 ONNX 模型

这里以 batchsize=1 为例。

batchsize=1
onnxsim models/encoder.onnx models/encoder_sim.onnx --overwrite-input-shape="x:${batchsize},3,1024,1024"
onnxsim models/decoder.onnx models/decoder_sim.onnx

参数说明：

第一个参数：原 ONNX 路径。
第二个参数：简化后的 ONNX 保存路径。
overwrite-input-shape：指定输入的维度。

5.4 运行改图脚本，修改 ONNX 模型以适配昇腾芯片

python3 encoder_onnx_modify.py \
  --input models/encoder_sim.onnx \
  --output models/encoder_modify.onnx

参数说明：

第一个参数：原 ONNX 路径。
第二个参数：适配后的 ONNX 保存路径。

5.5 使用 ATC 工具将 ONNX 模型转为 OM 模型

配置环境变量。
```
source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/mindie/set_env.sh
```
说明： 该脚本中环境变量仅供参考，请以实际安装环境配置环境变量。详细介绍请参见《CANN 开发辅助工具指南 $推理$ 》。

执行命令查看芯片名称（${chip_name}）。

npu-smi info
#该设备芯片名为Ascend310P3 （自行替换）
回显如下：
+-------------------+-----------------+------------------------------------------------------+
| NPU     Name      | Health          | Power(W)     Temp(C)           Hugepages-Usage(page) |
| Chip    Device    | Bus-Id          | AICore(%)    Memory-Usage(MB)                        |
+===================+=================+======================================================+
| 0       310P3     | OK              | 15.8         42                0    / 0              |
| 0       0         | 0000:82:00.0    | 0            1074 / 21534                            |
+===================+=================+======================================================+
| 1       310P3     | OK              | 15.4         43                0    / 0              |
| 0       1         | 0000:89:00.0    | 0            1070 / 21534                            |
+===================+=================+======================================================+

执行 atc 命令。
```
atc \
  --framework 5 \
  --model models/encoder_modify.onnx \
  --output models/encoder_modify \
  --insert_op_conf sam.aippconfig \
  --enable_small_channel 1 \
  --soc_version "Ascend${chip_name}"
atc \
  --framework 5 \
  --model models/decoder_sim.onnx \
  --output models/decoder_sim \
  --input_format 'ND' \
  --input_shape 'image_embeddings:1,256,64,64;point_coords:1,-1,2;point_labels:1,-1;mask_input:1,1,256,256;has_mask_input:1' \
  --dynamic_dims '2,2;3,3;4,4;5,5;6,6;7,7;8,8;9,9' \
  --op_select_implmode 'high_performance' \
  --optypelist_for_implmode 'Gelu' \
  --soc_version "Ascend${chip_name}"
```
参数说明：
- framework：原始框架类型。
- model：原始模型文件路径与文件名。
- output：存放转换后的离线模型的路径以及文件名。
- insert_op_conf：插入算子的配置文件路径与文件名，例如aipp预处理算子。
- enable_small_channel：是否使能small channel的优化，使能后在channel<=4的首层卷积会有性能收益。
- input_format：输入数据格式。
- input_shape：模型输入数据的shape。
- dynamic_dims：设置ND格式下动态维度的档位。适用于执行推理时，每次处理任意维度的场景。
- soc_version：模型转换时指定芯片版本。
更多参数说明请参考 ATC 参数概览（如果链接失效，请从 CANN 社区版文档查找目录「应用开发 > ATC 模型转换 > 参数说明 > 参数概览」）

6 推理验证

6.1 端到端推理。成功执行下述命令后会在save-path参数指定的目录生成离线推理的结果。

python3 sam_end2end_infer.py \
  --src-path data/demo.jpg \
  --save-path . \
  --encoder-model-path models/encoder_modify.om \
  --decoder-model-path models/decoder_sim.om \
  --input-point '[[500, 375], [1125, 625], [1520, 625]]' \
  --device-id 0

参数说明：

src-path：图片数据路径。
save-path：SAM离线推理保存路径。
encoder-model-path：encoder模型路径。
decoder-model-path：decoder模型路径。
input-point：分割目标上的坐标（坐标数量根据带分割目标物体分割效果确定，不同的图片数据，不同的待分割目标坐标不一样）
device-id：NPU卡ID

```

若出现类似"OSError: .../python3.10/site-packages/torch/lib/libgomp-d22c30c5.so.1: cannot allocate memory in static TLS block"错误，可以使用export LD_PRELOAD=$LD_PRELOAD:<报错路径>解决。
例子：export LD_PRELOAD=$LD_PRELOAD:/root/anaconda3/envs/sam/lib/python3.10/site-packages/torch/lib/libgomp-d22c30c5.so.1
```

在线模型推理结果：

离线模型推理结果：

6.2 性能验证。

encoder 纯推理性能验证。
```
python3 -m ais_bench --model models/encoder_modify.om --loop 100
```
参数说明：
- model: OM 模型。
- loop: 循环次数。

decoder 纯推理性能验证。

这里以输入 3 个坐标为例。

num_points=3
python3 -m ais_bench \
  --model models/decoder_sim.om \
  --dymDims "image_embeddings:1,256,64,64;point_coords:1,${num_points},2;point_labels:1,${num_points};mask_input:1,1,256,256;has_mask_input:1" \
  --outputSize '1000,1000000' \
  --loop 100

参数说明：

model: om模型
outputSize:动态模型输出Size设置
auto_set_dymdims_mode：开启动态dims模式
loop: 循环次数
batchsize: 模型batch size

6.3 精度验证。

SAM 官方未提供精度评测手段，这里提供对应脚本，基于 COCO 验证集标注框作为输入提示，使用 SAM 预测分割掩码，并与 COCO 标注掩码逐实例进行 IoU 计算，最后对所有实例的 IoU 结果取平均，得到整体的平均交并比（mIoU）。

```bash
python sam_coco_metric.py \
    --dataset-path coco2017 \
    --save-path outputs \
    --encoder-model-path models/encoder_sim.om \
    --decoder-model-path models/decoder_sim.om \
    --device-id 0 \
    --max-instances 0
```

参数说明：

dataset-path: coco数据集目录
save-path: SAM预测掩码结果存储路径
encoder-model-path：encoder的OM模型路径
decoder-model-path：decoder的OM模型路径
device-id: 指定推理的NPU设备ID
max-instances: 评测的最大实例数量，默认为0表示测评完整验证集

4. 模型推理性能 & 精度

性能结果：

芯片型号	模型	Batch Size	性能
300I Pro	encoder	1	4.43 fps
300I Pro	decoder	1	679.77 fps

精度结果：

芯片型号	模型	Batch Size	精度（mIoU）
300I Pro	SAM	1	0.7654

SAM(ONNX)-推理指导

概述

论文：

[Segment Anything](https://arxiv.org/abs/2304.02643)
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, Ross Girshick

参考实现：

url=https://github.com/facebookresearch/segment-anything.git
commit_id=6fdee8f2727f4506cfbbe553e23b895e27956588
model_name=sam_vit_b_01ec64

推理环境准备

该模型需要以下插件与驱动

表 1 版本配套表

配套	版本	环境准备指导
固件与驱动	25.2.RC1	Pytorch框架推理环境准备
CANN	8.2.RC1	-
Python	3.11.10	-
PyTorch	2.1.0	-
说明：仅支持Atlas 300I Duo 推理卡，请以CANN版本选择实际固件与驱动版本。	\	\

快速上手

1. 获取源码

git clone https://gitee.com/ascend/ModelZoo-PyTorch.git
cd ModelZoo-PyTorch/ACL_PyTorch/built-in/cv/SAM
git clone https://github.com/facebookresearch/segment-anything.git
cd segment-anything
git reset --hard 6fdee8f2727f4506cfbbe553e23b895e27956588
git apply ../segment_anything_diff.patch
pip3 install -e .
cd ..

2. 安装依赖。

安装基础环境。

pip3 install -r requirements.txt

说明：如果某些库通过此方式安装失败，可使用 pip3 install 单独进行安装。

安装 msit 的 surgeon 组件和 benchmark 组件。

3. 输入输出数据描述

encoder输入数据

输入数据数据类型大小数据排布格式
x FLOAT32 1 x 3 x 1024 x 1024 NCHW
encoder输出数据

输出数据数据类型大小数据排布格式
image_embeddings FLOAT32 1 x 256 x 64 x 64 NCHW

输入数据	数据类型	大小	数据排布格式
x	FLOAT32	1 x 3 x 1024 x 1024	NCHW

decoder输入数据

输入数据	数据类型	大小	数据排布格式
image_embeddings	FLOAT32	1 x 256 x 64 x 64	NCHW
point_coords	FLOAT32	1 x -1 x 2	ND
point_labels	FLOAT32	1 x -1	ND
mask_input	FLOAT32	1 x 1 x 256 x 256	NCHW
has_mask_input	FLOAT32	1	ND

decoder输出数据

输出数据数据类型大小数据排布格式
iou_predictions FLOAT32 -1 x 1 ND
low_res_masks FLOAT32 -1 x 1 x -1 x -1 ND

输出数据	数据类型	大小	数据排布格式
iou_predictions	FLOAT32	-1 x 1	ND
low_res_masks	FLOAT32	-1 x 1 x -1 x -1	ND

4. 准备数据集

取仓库里的 demo 图片进行端到端测试。

mkdir data
cd data
wget -O demo.jpg https://raw.githubusercontent.com/facebookresearch/segment-anything/571794162e0887c15d12b809505b902c7bf8b4db/notebooks/images/truck.jpg
cd ..

下载coco2017数据集进行精度测试。

下载并解压COCO-2017数据集的图片与标注，放置coco2017目录下

coco2017
├── annotations/
│   └── instances_val2017.json
└── val2017/
    ├── 000000000139.jpg
    ├── 000000000139.jpg
    └── ...

5. 模型转换

5.1 获取权重文件

GitHub 仓库提供了三种大小的权重文件：vit_h、vit_l、vit_b。这里以 vit_b 为例。

mkdir models
cd models
wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
cd ..

5.2 导出 ONNX 模型

python3 segment-anything/scripts/export_onnx_model.py  \
  --checkpoint models/sam_vit_b_01ec64.pth  \
  --model-type vit_b  \
  --opset 14 \
  --encoder-output models/encoder.onnx \
  --decoder-output models/decoder.onnx \
  --return-single-mask

参数说明：

checkpoint：模型权重文件路径。
model-type：模型类型。
opset：onnx算子集版本。
encoder-output：保存encoder模型的输出ONNX模型的文件路径。
decoder-output：保存decoder模型的输出ONNX模型的文件路径。
return-single-mask：设置最优mask模式。

5.3 使用 onnxsim 简化 ONNX 模型

这里以 batchsize=1 为例。

batchsize=1
onnxsim models/encoder.onnx models/encoder_sim.onnx --overwrite-input-shape="x:${batchsize},3,1024,1024"
onnxsim models/decoder.onnx models/decoder_sim.onnx

参数说明：

第一个参数：原 ONNX 路径。
第二个参数：简化后的 ONNX 保存路径。
overwrite-input-shape：指定输入的维度。

5.4 运行改图脚本，修改 ONNX 模型以适配昇腾芯片

python3 encoder_onnx_modify.py \
  --input models/encoder_sim.onnx \
  --output models/encoder_modify.onnx

参数说明：

第一个参数：原 ONNX 路径。
第二个参数：适配后的 ONNX 保存路径。

5.5 使用 ATC 工具将 ONNX 模型转为 OM 模型

配置环境变量。
```
source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/mindie/set_env.sh
```
说明： 该脚本中环境变量仅供参考，请以实际安装环境配置环境变量。详细介绍请参见《CANN 开发辅助工具指南 $推理$ 》。

执行命令查看芯片名称（${chip_name}）。

npu-smi info
#该设备芯片名为Ascend310P3 （自行替换）
回显如下：
+-------------------+-----------------+------------------------------------------------------+
| NPU     Name      | Health          | Power(W)     Temp(C)           Hugepages-Usage(page) |
| Chip    Device    | Bus-Id          | AICore(%)    Memory-Usage(MB)                        |
+===================+=================+======================================================+
| 0       310P3     | OK              | 15.8         42                0    / 0              |
| 0       0         | 0000:82:00.0    | 0            1074 / 21534                            |
+===================+=================+======================================================+
| 1       310P3     | OK              | 15.4         43                0    / 0              |
| 0       1         | 0000:89:00.0    | 0            1070 / 21534                            |
+===================+=================+======================================================+

执行 atc 命令。
```
atc \
  --framework 5 \
  --model models/encoder_modify.onnx \
  --output models/encoder_modify \
  --insert_op_conf sam.aippconfig \
  --enable_small_channel 1 \
  --soc_version "Ascend${chip_name}"
atc \
  --framework 5 \
  --model models/decoder_sim.onnx \
  --output models/decoder_sim \
  --input_format 'ND' \
  --input_shape 'image_embeddings:1,256,64,64;point_coords:1,-1,2;point_labels:1,-1;mask_input:1,1,256,256;has_mask_input:1' \
  --dynamic_dims '2,2;3,3;4,4;5,5;6,6;7,7;8,8;9,9' \
  --op_select_implmode 'high_performance' \
  --optypelist_for_implmode 'Gelu' \
  --soc_version "Ascend${chip_name}"
```
参数说明：
- framework：原始框架类型。
- model：原始模型文件路径与文件名。
- output：存放转换后的离线模型的路径以及文件名。
- insert_op_conf：插入算子的配置文件路径与文件名，例如aipp预处理算子。
- enable_small_channel：是否使能small channel的优化，使能后在channel<=4的首层卷积会有性能收益。
- input_format：输入数据格式。
- input_shape：模型输入数据的shape。
- dynamic_dims：设置ND格式下动态维度的档位。适用于执行推理时，每次处理任意维度的场景。
- soc_version：模型转换时指定芯片版本。
更多参数说明请参考 ATC 参数概览（如果链接失效，请从 CANN 社区版文档查找目录「应用开发 > ATC 模型转换 > 参数说明 > 参数概览」）

6 推理验证

6.1 端到端推理。成功执行下述命令后会在save-path参数指定的目录生成离线推理的结果。

python3 sam_end2end_infer.py \
  --src-path data/demo.jpg \
  --save-path . \
  --encoder-model-path models/encoder_modify.om \
  --decoder-model-path models/decoder_sim.om \
  --input-point '[[500, 375], [1125, 625], [1520, 625]]' \
  --device-id 0

参数说明：

src-path：图片数据路径。
save-path：SAM离线推理保存路径。
encoder-model-path：encoder模型路径。
decoder-model-path：decoder模型路径。
input-point：分割目标上的坐标（坐标数量根据带分割目标物体分割效果确定，不同的图片数据，不同的待分割目标坐标不一样）
device-id：NPU卡ID

```

若出现类似"OSError: .../python3.10/site-packages/torch/lib/libgomp-d22c30c5.so.1: cannot allocate memory in static TLS block"错误，可以使用export LD_PRELOAD=$LD_PRELOAD:<报错路径>解决。
例子：export LD_PRELOAD=$LD_PRELOAD:/root/anaconda3/envs/sam/lib/python3.10/site-packages/torch/lib/libgomp-d22c30c5.so.1
```

在线模型推理结果：

离线模型推理结果：

6.2 性能验证。

encoder 纯推理性能验证。
```
python3 -m ais_bench --model models/encoder_modify.om --loop 100
```
参数说明：
- model: OM 模型。
- loop: 循环次数。

decoder 纯推理性能验证。

这里以输入 3 个坐标为例。

num_points=3
python3 -m ais_bench \
  --model models/decoder_sim.om \
  --dymDims "image_embeddings:1,256,64,64;point_coords:1,${num_points},2;point_labels:1,${num_points};mask_input:1,1,256,256;has_mask_input:1" \
  --outputSize '1000,1000000' \
  --loop 100

参数说明：

model: om模型
outputSize:动态模型输出Size设置
auto_set_dymdims_mode：开启动态dims模式
loop: 循环次数
batchsize: 模型batch size

6.3 精度验证。

```bash
python sam_coco_metric.py \
    --dataset-path coco2017 \
    --save-path outputs \
    --encoder-model-path models/encoder_sim.om \
    --decoder-model-path models/decoder_sim.om \
    --device-id 0 \
    --max-instances 0
```

参数说明：

dataset-path: coco数据集目录
save-path: SAM预测掩码结果存储路径
encoder-model-path：encoder的OM模型路径
decoder-model-path：decoder的OM模型路径
device-id: 指定推理的NPU设备ID
max-instances: 评测的最大实例数量，默认为0表示测评完整验证集

4. 模型推理性能 & 精度

性能结果：

芯片型号	模型	Batch Size	性能
300I Pro	encoder	1	4.43 fps
300I Pro	decoder	1	679.77 fps

精度结果：

芯片型号	模型	Batch Size	精度（mIoU）
300I Pro	SAM	1	0.7654