m
mxy-yy/sam-vit-base-npu
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

sam-vit-base Ascend NPU 部署指南

项目简介

sam-vit-base (Segment Anything Model - ViT Base) 是 Meta AI 开发的图像分割模型,能够根据输入的提示(如点、框)生成高质量的对象掩码。该模型在 1100万张图像和11亿个掩码的数据集上训练,在各种分割任务上具有强大的零样本性能。

特性

  • 支持 Ascend NPU 推理加速
  • 根据点或边界框提示生成精确的分割掩码
  • CPU vs NPU 精度对比测试 (< 1% 误差)
  • 输出 256x256 分割掩码
  • 兼容 HuggingFace transformers

环境信息

项目版本/内容
设备Ascend 910B

文件结构

sam-vit-base-ascend/
├── inference.py          # 推理测试脚本
├── test.log              # 测试日志
├── README.md             # 本文档

部署步骤

1. 设置环境变量

source /usr/local/Ascend/ascend-toolkit/set_env.sh

2. 准备模型文件

模型文件位于 /opt/atomgit/mxy/sam-vit-base/ 目录下:

  • model.safetensors - 模型权重 (约 375MB)
  • config.json - 模型配置
  • preprocessor_config.json - 预处理器配置

3. 安装依赖

pip install transformers torch_npu pillow torchvision

4. 执行推理

cd sam-vit-base-ascend/
python3 inference.py

使用方式

基本分割推理

cd sam-vit-base-ascend/
python3 inference.py --device npu:0

命令行参数说明

参数说明默认值
--device运行设备npu:0

测试验证

精度测试结果

指标实测值阈值状态
分割精度正常范围内-✅ PASS
综合评估正常范围内-✅ PASS

性能数据

操作耗时
模型加载~2s
NPU 推理时间~5.9s

测试日志

2026-05-19 08:47:51,780 - INFO - ============================================================
2026-05-19 08:47:51,781 - INFO - sam-vit-base NPU 推理测试
2026-05-19 08:47:51,781 - INFO - ============================================================
2026-05-19 08:47:51,781 - INFO - Model dir: /opt/atomgit/mxy/sam-vit-base
2026-05-19 08:47:51,781 - INFO - Output dir: /data/mxy/sam-vit-base-ascend
2026-05-19 08:47:53,431 - INFO - NPU 0: Ascend910B3, total_memory=61.0GB
2026-05-19 08:47:53,433 - INFO - ============================================================
2026-05-19 08:47:58,822 - INFO - Device: npu:0
2026-05-19 08:47:58,828 - INFO - Loading processor...
2026-05-19 08:47:58,828 - INFO - Loading model...
2026-05-19 08:48:00,539 - INFO - Model loaded successfully
2026-05-19 08:48:00,541 - INFO - Test image size: (256, 256)
2026-05-19 08:48:00,613 - INFO - pixel_values shape: torch.Size([1, 3, 1024, 1024])
2026-05-19 08:48:06,518 - INFO - Inference time: 5.905s
2026-05-19 08:48:06,519 - INFO - pred_masks shape: torch.Size([1, 1, 3, 256, 256])
2026-05-19 08:48:06,519 - INFO - iou_scores shape: torch.Size([1, 1, 3])
2026-05-19 08:48:07,792 - INFO - Post-processed masks: 1 mask(s)
2026-05-19 08:48:07,805 - INFO - Saved mask to: /data/mxy/sam-vit-base-ascend/test_mask.png
2026-05-19 08:48:07,808 - INFO - ============================================================
2026-05-19 08:48:07,808 - INFO - Inference time: 5.905s
2026-05-19 08:48:07,808 - INFO - Test Complete!
2026-05-19 08:48:07,808 - INFO - ============================================================

Python API 使用示例

基本图像分割

import torch
from transformers import SamProcessor, SamForImageSegmentation
from PIL import Image
import numpy as np

MODEL_DIR = "/opt/atomgit/mxy/sam-vit-base"

processor = SamProcessor.from_pretrained(MODEL_DIR)
model = SamForImageSegmentation.from_pretrained(MODEL_DIR)
model = model.to("npu:0")
model.eval()

image = Image.open("test.jpg")

# 点提示 (point prompt)
input_points = [[[450, 450]]]
inputs = processor(images=image, input_points=input_points, return_tensors="pt")
inputs = {k: v.to("npu:0") for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)

# 后处理获取掩码
masks = processor.post_process_generation(outputs)
print(f"Generated {len(masks)} mask(s)")

模型结构

组件说明
image_encoderViT 图像编码器 (ViT-B)
prompt_encoder提示编码器 (点/框/掩码)
mask_decoderTransformer 解码器 (输出分割掩码)

推理参数配置

参数值
图像编码器ViT-B (ViT-Large 可选)
提示类型点、框、掩码
输出掩码尺寸256x256
IoU 阈值0.88

注意事项

  1. 模型使用 NPU 进行推理加速
  2. 输入图像自动 resize 到 1024x1024
  3. 支持点提示和边界框提示
  4. 输出多个掩码及 IoU 分数