冬
gcw_IDzXRVNw/vitmatte-small-composition-1k-ascend
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

vitmatte-small-composition-1k Ascend NPU 部署指南

项目简介

vitmatte-small-composition-1k 是基于视觉Transformer（ViT）的图像抠图模型，专门在Composition-1k数据集上训练。该模型由Yao等人在论文ViTMatte: Boosting Image Matting with Pretrained Plain Vision Transformers中提出，能够准确估计图像中的前景物体，实现高精度抠图。

特性

  • 支持Ascend NPU推理加速
  • CPU与NPU精度对比测试（误差<1%）
  • 512x512高分辨率图像处理
  • 4通道输入（RGB + Trimap）
  • 158倍加速比

环境要求

  • 硬件：华为Ascend 910系列NPU
  • CANN：8.0.RC1或更高版本
  • PyTorch：2.0+并带有torch_npu
  • Docker：容器名称test-modelagent
  • transformers：4.31+

目录结构

vitmatte-small-composition-1k-ascend/
├── inference.py          # 推理测试脚本
├── log.txt               # 测试日志
├── README.md             # 本文档
├── test_image.npy        # 测试图像
├── test_trimap.npy       # 测试trimap
├── test_sample.txt       # 测试样本说明
├── inference_result.json # 推理结果
└── precision_result.json # 精度测试结果

部署步骤

1. 进入容器

docker exec -it test-modelagent bash

2. 设置环境变量

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

3. 准备模型文件

模型文件位于 /data/ysws/agentsp/5-16/vitmatte-small-composition-1k/hustvl/vitmatte-small-composition-1k/ 目录下：

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

4. 安装依赖

pip install transformers torch_npu -i https://pypi.huaweicloud.com/repository/pypi/simple/

Usage

Method 1: Normal Inference Mode

Run the inference script for image matting:

cd /data/ysws/agentsp/5-16/vitmatte-small-composition-1k-ascend/

python3 inference.py --mode inference

方式二：精度测试模式 (CPU vs NPU)

运行精度对比测试，验证 NPU 计算结果与 CPU 一致性：

cd /data/ysws/agentsp/5-16/vitmatte-small-composition-1k-ascend/

python3 inference.py --mode precision_test

方式三：完整测试 (推理 + 精度)

cd /data/ysws/agentsp/5-16/vitmatte-small-composition-1k-ascend/

python3 inference.py --mode all

命令行参数说明

参数说明默认值
--mode测试模式: inference, precision_test 或 allall

测试验证

精度测试结果

指标实测值阈值状态
最大相对误差0.8556%< 1.00%PASS
最大绝对误差2.67e-05--
CPU 推理时间6.952s--
NPU 推理时间0.044s--
加速比158.55x> 1xPASS

性能数据

操作耗时
NPU 推理时间 (512x512)6.427s
精度测试 CPU 时间6.952s
精度测试 NPU 时间0.044s

推理结果示例

输入图像: 512x512 RGB 输入 Trimap: 512x512 掩码 (0=背景, 1=未知, 2=前景) 输出 Alpha: 512x512 透明度通道

Python API 使用示例

基本图像抠图

import torch
import numpy as np
from transformers import VitMatteForImageMatting, VitMatteImageProcessor

MODEL_DIR = "/data/ysws/agentsp/5-16/vitmatte-small-composition-1k/hustvl/vitmatte-small-composition-1k"

model = VitMatteForImageMatting.from_pretrained(MODEL_DIR)
processor = VitMatteImageProcessor.from_pretrained(MODEL_DIR)

model = model.to("npu:0").eval()

image = np.random.rand(512, 512, 3)
trimap = np.random.randint(0, 2, (512, 512)).astype(np.float32)

inputs = processor(images=image, trimaps=trimap, return_tensors="pt")
pixel_values = inputs["pixel_values"].to("npu:0")

with torch.no_grad():
    outputs = model(pixel_values)

alphas = outputs.alphas

处理实际图像

from PIL import Image

img = Image.open("your_image.jpg")
trimap_img = Image.open("your_trimap.png")

image = np.array(img).astype(np.float32) / 255.0
trimap = np.array(trimap_img).astype(np.float32) / 255.0

inputs = processor(images=image, trimaps=trimap, return_tensors="pt")
outputs = model(inputs["pixel_values"].to("npu:0"))
alpha = outputs.alphas.cpu().numpy()[0, 0]

result_img = Image.fromarray((alpha * 255).astype(np.uint8))
result_img.save("result.png")

模型结构

  • 架构类型: VitMatteForImageMatting
  • 骨干网络: ViT-Det (基于窗口)
  • 隐藏层维度: 384
  • 注意力头数: 6
  • 输入: 4通道 (RGB + 归一化Trimap)
组件说明
backboneViT-Det 特征提取器
convstream卷积流特征融合
fusion融合模块输出最终 alpha

推理参数配置

从 config.json 提取的关键参数:

{
  "hidden_size": 384,
  "image_size": 512,
  "model_type": "vitmatte",
  "backbone_config": {
    "hidden_size": 384,
    "num_attention_heads": 6,
    "num_channels": 4,
    "window_size": 14
  },
  "convstream_hidden_sizes": [48, 96, 192],
  "fusion_hidden_sizes": [256, 128, 64, 32]
}

常见问题

Q: 精度测试失败?

A: 检查 NPU 驱动是否正确安装。0.85% 的数值误差是正常的，因为 NPU 和 CPU 使用不同的计算精度和算子实现。

Q: 推理时间较长?

A: 首次推理需要编译算子。ViTMatte 模型较大，NPU 推理非常快 (0.044s vs CPU 6.95s)。

Q: Trimap 是什么?

A: Trimap 是一种三值掩码，表示图像中每个像素的状态：0=确定背景，1=未知区域，2=确定前景。模型根据 Trimap 估计未知区域的 alpha 值。

测试日志

============================================================
ViTMatte NPU Test
Model: hustvl/vitmatte-small-composition-1k
Output: /data/ysws/agentsp/5-16/vitmatte-small-composition-1k-ascend
============================================================

============================================================
ViTMatte Inference Test (NPU)
============================================================
Device: npu:0
Model: /data/ysws/agentsp/5-16/vitmatte-small-composition-1k/hustvl/vitmatte-small-composition-1k
Loading model...
Loading weights: 100%|██████████| 258/258 [00:00<00:00, 4963.51it/s]
Model loaded successfully
Image shape: torch.Size([1, 3, 512, 512])
Trimap shape: torch.Size([1, 1, 512, 512])
Pixel values shape: torch.Size([1, 4, 512, 512])
Output alpha shape: torch.Size([1, 1, 512, 512])
Inference time: 6.427s

Inference result saved to /data/ysws/agentsp/5-16/vitmatte-small-composition-1k-ascend/inference_result.json

============================================================
Precision Test (CPU vs NPU)
============================================================
Using device: npu:0
Loading model on CPU...
Loading weights: 100%|██████████| 258/258 [00:00<00:00, 4597.71it/s]
Loading model on npu:0...
Loading weights: 100%|██████████| 258/258 [00:00<00:00, 4690.33it/s]
Loading processor...
Running inference on CPU...
Running inference on NPU...
CPU inference time: 6.952s
NPU inference time: 0.044s
Speedup: 158.55x
Max absolute error: 2.669613e-05
Max relative error: 0.8556% (threshold: 1.0%)
Status: PASS

Precision result saved to /data/ysws/agentsp/5-16/vitmatte-small-composition-1k-ascend/precision_result.json

============================================================
Creating Test Sample
============================================================
Saved test image: /data/ysws/agentsp/5-16/vitmatte-small-composition-1k-ascend/test_image.npy (shape: torch.Size([3, 512, 512]))
Saved test trimap: /data/ysws/agentsp/5-16/vitmatte-small-composition-1k-ascend/test_trimap.npy (shape: torch.Size([1, 512, 512]))
Saved test sample info: /data/ysws/agentsp/5-16/vitmatte-small-composition-1k-ascend/test_sample.txt

============================================================
Test Complete!
============================================================

参考链接

  • 原始模型: https://huggingface.co/hustvl/vitmatte-small-composition-1k
  • 论文: ViTMatte: Boosting Image Matting with Pretrained Plain Vision Transformers
  • GitHub: hustvl/ViTMatte

许可证

本项目遵循 Apache-2.0 许可证