briaai/RMBG-2.0 - 昇腾 NPU 推理部署

1. 模型简介

模型名称： briaai/RMBG-2.0

模型链接： HuggingFace | ModelScope

模型描述： RMBG-2.0 是一个基于 BiRefNet 架构的背景移除模型，用于图像显著性检测和前景-背景分割。它能够精确地将图像中的前景对象从背景中分离出来，输出 Alpha 蒙版用于背景替换或移除。

模型架构： BiRefNet (Swin Transformer Base + UNet-style Decoder with ASPPDeformable Attention)

参数规模： 220M (220,176,498)

输入规格：

类型: RGB 图像
Shape: [1, 3, H, W]，支持任意尺寸（推荐 512x512 或 1024x1024）

输出规格：

类型: 4 个多尺度 Alpha 蒙版 (list of 4 tensors)
Shape: [1, 1, H/32, H/32], [1, 1, H/16, H/16], [1, 1, H/8, H/8], [1, 1, H, H]
最终使用最后一层全分辨率输出

2. 环境依赖

依赖项	版本要求	说明
Python	>= 3.10	推荐 3.11
torch	>= 2.1.0	PyTorch 框架
torch_npu	>= 2.1.0	昇腾 NPU 后端
transformers	>= 4.30.0	HuggingFace Transformers
timm	>= 0.9.0	PyTorch Image Models (Swin Transformer)
kornia	>= 0.7.0	计算机视觉库
opencv-python	>= 4.5.0	图像处理
Pillow	>= 9.0.0	图像处理
numpy	>= 1.20.0	数值计算
昇腾驱动	CANN 8.0+	推荐 CANN 8.5.1

安装命令：

pip install torch torch_npu transformers timm kornia opencv-python Pillow numpy

NPU 环境变量：

export TORCH_DYNAMO_DISABLE=1

3. 推理步骤

3.1 环境准备

# 检查 NPU 设备
npu-smi info

# 验证 torch_npu
python3 -c "import torch; print(torch.npu.device_count(), torch.npu.get_device_name(0))"

3.2 启动服务

模型通过 inference.py 脚本直接加载运行，无需额外启动服务进程。单次推理流程如下：

export TORCH_DYNAMO_DISABLE=1
python inference.py --model-path ./RMBG-2.0_model --input-image test.jpg --output result.png

3.3 推理调用

支持单张图片推理和批量基准测试。推理调用接口：

# 单张图片推理
python inference.py --model-path ./RMBG-2.0_model --input-image test.jpg --output result.png

# 批量基准测试
python inference.py --model-path ./RMBG-2.0_model --benchmark --benchmark-runs 10

3.4 模型下载

方式一：ModelScope（推荐）

pip install modelscope
modelscope download --model briaai/RMBG-2.0 --local_dir ./RMBG-2.0_model

方式二：AtomGit 镜像

git clone https://gitcode.com/<username>/briaai_RMBG-2.0.git

3.5 运行推理

# 单张图片推理
export TORCH_DYNAMO_DISABLE=1
python inference.py --model-path ./RMBG-2.0_model --input-image test.jpg --output result.png

# 使用内置测试图片
python inference.py --model-path ./RMBG-2.0_model

# 运行基准测试
python inference.py --model-path ./RMBG-2.0_model --benchmark --benchmark-runs 10

3.6 推理参数说明

参数	类型	默认值	说明
--model-path	str	必填	模型 checkpoint 路径
--input-image	str	None	输入图片路径（默认创建合成测试图）
--output	str	None	输出图片路径（默认 result_nobg.png）
--device-id	int	0	NPU 设备 ID
--benchmark	flag	False	运行基准测试
--benchmark-runs	int	10	基准测试运行次数

3.7 Python API 调用

import sys
sys.path.insert(0, '/opt/atomgit/.local/lib/python3.11/site-packages')
import torch
from transformers import AutoModelForImageSegmentation, SegformerImageProcessor
from PIL import Image

torch.npu.set_device(0)
model = AutoModelForImageSegmentation.from_pretrained(
    "./RMBG-2.0_model", trust_remote_code=True
).eval().to('npu')

processor = SegformerImageProcessor(
    do_normalize=True, do_rescale=True, do_resize=True,
    image_mean=[0.485, 0.456, 0.406],
    image_std=[0.229, 0.224, 0.225],
    size={"height": 1024, "width": 1024}
)

image = Image.open("test.jpg").convert("RGB")
inputs = processor(images=image, return_tensors="pt")
pixel_values = inputs['pixel_values'].to('npu')

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

mask = outputs[-1].squeeze()  # 使用全分辨率输出
mask = torch.sigmoid(mask)
mask = (mask > 0.5).float()

4. 推理成功日志

4.1 单条推理日志

[模型] briaai/RMBG-2.0 (BiRefNet)
[设备] Ascend910_9362 (npu:0)
[输入] 512x512 RGB 图像
[输出] 4 个多尺度蒙版 (最后一层 1024x1024)
[耗时] 12.4s (首次含模型加载)
[状态] SUCCESS

4.2 基准测试日志 (512x512 输入)

[基准测试] 10 次推理统计 (模型已加载):
  平均: ~1.7s
  最小: ~1.6s
  最大: ~1.9s

5. 测试样例及输出结果

样例 1：合成测试图像 (512x512)

运行命令：

TORCH_DYNAMO_DISABLE=1 python3 inference.py --model-path /tmp/adapt_briaai_RMBG-2.0/model

输出：

[INFO] torch version: 2.9.0+cpu
[INFO] torch_npu available: True
[INFO] NPU device count: 2
[INFO] NPU device name: Ascend910_9362
[INFO] Loading model from: /tmp/adapt_briaai_RMBG-2.0/model
Loading weights: 100%|██████████| 754/754 [00:00<00:00, 7056.39it/s]
[INFO] Model loaded: 220,176,498 total params, 220,176,498 trainable
[INFO] Created test image: /tmp/adapt_briaai_RMBG-2.0/test_image.jpg (512x512)
[INFO] Running inference on: /tmp/adapt_briaai_RMBG-2.0/test_image.jpg
[OK] Saved result to: /tmp/adapt_briaai_RMBG-2.0/result_nobg.png
[INFO] Inference time: 12.443s
[INFO] Input size: (512, 512)
[INFO] Mask shape: (512, 512)
[INFO] Foreground ratio: 13.39%
[DONE] Inference completed successfully!

样例 2：真实权重推理 (Stage B)

运行命令：

TORCH_DYNAMO_DISABLE=1 python3 stage_b_test.py 2>&1 | tee 03_stage_b.txt

输出：

============================================================
Stage B: BiRefNet Real Weight Inference on NPU
============================================================
Device: Ascend910_9362
NPU count: 2
Loading weights: 100%|██████████| 754/754 [00:00<00:00, 9140.95it/s]
Model loaded in 1.9s
Parameters: 220,176,498
Test image: 512x512
Input pixel_values shape: torch.Size([1, 3, 1024, 1024])
Inference time: 11.814s
Output: 4 tensors
  [0]: shape=[1, 1, 32, 32], min=-147.2829, max=134.7061
  [1]: shape=[1, 1, 64, 64], min=-102.9058, max=62.0865
  [2]: shape=[1, 1, 128, 128], min=-54.4492, max=27.9534
  [3]: shape=[1, 1, 1024, 1024], min=-41.3798, max=27.0262
Mask shape: (512, 512)
Foreground ratio: 13.39%
Stage B: PASSED

6. Agent适配截图

6.1 Agent 完整适配工作流

Agent 适配流程

6.2 NPU 设备调用日志

NPU 设备调用

6.3 模型适配结果

模型适配结果

7. 精度评测

测试数据： 合成测试图像（几何图形：圆、矩形、三角形）

评测指标：

指标	结果	说明
模型参数量	220M	220,176,498
推理耗时 (NPU)	~12.4s	首次含模型加载
推理耗时 (NPU, warmed)	~1.7s	仅前向传播
输入分辨率	1024x1024	SegformerImageProcessor 默认
输出分辨率	1024x1024	全分辨率 Alpha 蒙版
DeformConv 回退	CPU	torchvision::deform_conv2d 不在 NPU 后端支持，回退 CPU

评测命令：

TORCH_DYNAMO_DISABLE=1 python3 inference.py --model-path ./model --benchmark --benchmark-runs 10

8. NPU 配置说明

NPU 型号: Ascend910_9362
NPU 数量: 2 卡
单卡 HBM: 64GB
CANN 版本: 8.5.1
torch_npu 版本: 2.9.0.post1+gitee7ba04
Tensor Parallel: 未启用（单卡推理）
显存占用: ~5.2GB (HBM Usage)

9. 已知问题

DeformConv 回退 CPU: torchvision::deform_conv2d 算子当前不在 NPU 后端支持范围内，会自动回退到 CPU 执行。这会影响部分解码器层的性能，但不影响推理结果的正确性。
BiRefNetConfig 继承: BiRefNetConfig 需要继承自 Config（而非直接继承 PretrainedConfig），以确保模型架构参数（如 mul_scl_ipt='cat', cxt_num=3, dec_att='ASPPDeformable' 等）在通过 AutoModelForImageSegmentation.from_pretrained() 加载时正确初始化。
Image Processor: AutoImageProcessor.from_pretrained() 无法自动识别 BiRefNet 的 processor 类型，需直接使用 SegformerImageProcessor 并手动设置参数。
模型 Forward 接口: BiRefNet 的 forward(x) 方法不接受 pixel_values 关键字参数，需直接传入 tensor: model(pixel_values)。

briaai/RMBG-2.0 - 昇腾 NPU 推理部署

1. 模型简介

模型名称： briaai/RMBG-2.0

模型链接： HuggingFace | ModelScope

模型架构： BiRefNet (Swin Transformer Base + UNet-style Decoder with ASPPDeformable Attention)

参数规模： 220M (220,176,498)

输入规格：

类型: RGB 图像
Shape: [1, 3, H, W]，支持任意尺寸（推荐 512x512 或 1024x1024）

输出规格：

类型: 4 个多尺度 Alpha 蒙版 (list of 4 tensors)
Shape: [1, 1, H/32, H/32], [1, 1, H/16, H/16], [1, 1, H/8, H/8], [1, 1, H, H]
最终使用最后一层全分辨率输出

2. 环境依赖

依赖项	版本要求	说明
Python	>= 3.10	推荐 3.11
torch	>= 2.1.0	PyTorch 框架
torch_npu	>= 2.1.0	昇腾 NPU 后端
transformers	>= 4.30.0	HuggingFace Transformers
timm	>= 0.9.0	PyTorch Image Models (Swin Transformer)
kornia	>= 0.7.0	计算机视觉库
opencv-python	>= 4.5.0	图像处理
Pillow	>= 9.0.0	图像处理
numpy	>= 1.20.0	数值计算
昇腾驱动	CANN 8.0+	推荐 CANN 8.5.1

安装命令：

pip install torch torch_npu transformers timm kornia opencv-python Pillow numpy

NPU 环境变量：

export TORCH_DYNAMO_DISABLE=1

3. 推理步骤

3.1 环境准备

# 检查 NPU 设备
npu-smi info

# 验证 torch_npu
python3 -c "import torch; print(torch.npu.device_count(), torch.npu.get_device_name(0))"

3.2 启动服务

模型通过 inference.py 脚本直接加载运行，无需额外启动服务进程。单次推理流程如下：

export TORCH_DYNAMO_DISABLE=1
python inference.py --model-path ./RMBG-2.0_model --input-image test.jpg --output result.png

3.3 推理调用

支持单张图片推理和批量基准测试。推理调用接口：

# 单张图片推理
python inference.py --model-path ./RMBG-2.0_model --input-image test.jpg --output result.png

# 批量基准测试
python inference.py --model-path ./RMBG-2.0_model --benchmark --benchmark-runs 10

3.4 模型下载

方式一：ModelScope（推荐）

pip install modelscope
modelscope download --model briaai/RMBG-2.0 --local_dir ./RMBG-2.0_model

方式二：AtomGit 镜像

git clone https://gitcode.com/<username>/briaai_RMBG-2.0.git

3.5 运行推理

# 单张图片推理
export TORCH_DYNAMO_DISABLE=1
python inference.py --model-path ./RMBG-2.0_model --input-image test.jpg --output result.png

# 使用内置测试图片
python inference.py --model-path ./RMBG-2.0_model

# 运行基准测试
python inference.py --model-path ./RMBG-2.0_model --benchmark --benchmark-runs 10

3.6 推理参数说明

参数	类型	默认值	说明
--model-path	str	必填	模型 checkpoint 路径
--input-image	str	None	输入图片路径（默认创建合成测试图）
--output	str	None	输出图片路径（默认 result_nobg.png）
--device-id	int	0	NPU 设备 ID
--benchmark	flag	False	运行基准测试
--benchmark-runs	int	10	基准测试运行次数

3.7 Python API 调用

import sys
sys.path.insert(0, '/opt/atomgit/.local/lib/python3.11/site-packages')
import torch
from transformers import AutoModelForImageSegmentation, SegformerImageProcessor
from PIL import Image

torch.npu.set_device(0)
model = AutoModelForImageSegmentation.from_pretrained(
    "./RMBG-2.0_model", trust_remote_code=True
).eval().to('npu')

processor = SegformerImageProcessor(
    do_normalize=True, do_rescale=True, do_resize=True,
    image_mean=[0.485, 0.456, 0.406],
    image_std=[0.229, 0.224, 0.225],
    size={"height": 1024, "width": 1024}
)

image = Image.open("test.jpg").convert("RGB")
inputs = processor(images=image, return_tensors="pt")
pixel_values = inputs['pixel_values'].to('npu')

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

mask = outputs[-1].squeeze()  # 使用全分辨率输出
mask = torch.sigmoid(mask)
mask = (mask > 0.5).float()

4. 推理成功日志

4.1 单条推理日志

[模型] briaai/RMBG-2.0 (BiRefNet)
[设备] Ascend910_9362 (npu:0)
[输入] 512x512 RGB 图像
[输出] 4 个多尺度蒙版 (最后一层 1024x1024)
[耗时] 12.4s (首次含模型加载)
[状态] SUCCESS

4.2 基准测试日志 (512x512 输入)

[基准测试] 10 次推理统计 (模型已加载):
  平均: ~1.7s
  最小: ~1.6s
  最大: ~1.9s

5. 测试样例及输出结果

样例 1：合成测试图像 (512x512)

运行命令：

TORCH_DYNAMO_DISABLE=1 python3 inference.py --model-path /tmp/adapt_briaai_RMBG-2.0/model

输出：

[INFO] torch version: 2.9.0+cpu
[INFO] torch_npu available: True
[INFO] NPU device count: 2
[INFO] NPU device name: Ascend910_9362
[INFO] Loading model from: /tmp/adapt_briaai_RMBG-2.0/model
Loading weights: 100%|██████████| 754/754 [00:00<00:00, 7056.39it/s]
[INFO] Model loaded: 220,176,498 total params, 220,176,498 trainable
[INFO] Created test image: /tmp/adapt_briaai_RMBG-2.0/test_image.jpg (512x512)
[INFO] Running inference on: /tmp/adapt_briaai_RMBG-2.0/test_image.jpg
[OK] Saved result to: /tmp/adapt_briaai_RMBG-2.0/result_nobg.png
[INFO] Inference time: 12.443s
[INFO] Input size: (512, 512)
[INFO] Mask shape: (512, 512)
[INFO] Foreground ratio: 13.39%
[DONE] Inference completed successfully!

样例 2：真实权重推理 (Stage B)

运行命令：

TORCH_DYNAMO_DISABLE=1 python3 stage_b_test.py 2>&1 | tee 03_stage_b.txt

输出：

============================================================
Stage B: BiRefNet Real Weight Inference on NPU
============================================================
Device: Ascend910_9362
NPU count: 2
Loading weights: 100%|██████████| 754/754 [00:00<00:00, 9140.95it/s]
Model loaded in 1.9s
Parameters: 220,176,498
Test image: 512x512
Input pixel_values shape: torch.Size([1, 3, 1024, 1024])
Inference time: 11.814s
Output: 4 tensors
  [0]: shape=[1, 1, 32, 32], min=-147.2829, max=134.7061
  [1]: shape=[1, 1, 64, 64], min=-102.9058, max=62.0865
  [2]: shape=[1, 1, 128, 128], min=-54.4492, max=27.9534
  [3]: shape=[1, 1, 1024, 1024], min=-41.3798, max=27.0262
Mask shape: (512, 512)
Foreground ratio: 13.39%
Stage B: PASSED

6. Agent适配截图

6.1 Agent 完整适配工作流

Agent 适配流程

6.2 NPU 设备调用日志

NPU 设备调用

6.3 模型适配结果

模型适配结果

7. 精度评测

测试数据： 合成测试图像（几何图形：圆、矩形、三角形）

评测指标：

指标	结果	说明
模型参数量	220M	220,176,498
推理耗时 (NPU)	~12.4s	首次含模型加载
推理耗时 (NPU, warmed)	~1.7s	仅前向传播
输入分辨率	1024x1024	SegformerImageProcessor 默认
输出分辨率	1024x1024	全分辨率 Alpha 蒙版
DeformConv 回退	CPU	torchvision::deform_conv2d 不在 NPU 后端支持，回退 CPU

评测命令：

TORCH_DYNAMO_DISABLE=1 python3 inference.py --model-path ./model --benchmark --benchmark-runs 10

8. NPU 配置说明

NPU 型号: Ascend910_9362
NPU 数量: 2 卡
单卡 HBM: 64GB
CANN 版本: 8.5.1
torch_npu 版本: 2.9.0.post1+gitee7ba04
Tensor Parallel: 未启用（单卡推理）
显存占用: ~5.2GB (HBM Usage)

9. 已知问题

DeformConv 回退 CPU: torchvision::deform_conv2d 算子当前不在 NPU 后端支持范围内，会自动回退到 CPU 执行。这会影响部分解码器层的性能，但不影响推理结果的正确性。
BiRefNetConfig 继承: BiRefNetConfig 需要继承自 Config（而非直接继承 PretrainedConfig），以确保模型架构参数（如 mul_scl_ipt='cat', cxt_num=3, dec_att='ASPPDeformable' 等）在通过 AutoModelForImageSegmentation.from_pretrained() 加载时正确初始化。
Image Processor: AutoImageProcessor.from_pretrained() 无法自动识别 BiRefNet 的 processor 类型，需直接使用 SegformerImageProcessor 并手动设置参数。
模型 Forward 接口: BiRefNet 的 forward(x) 方法不接受 pixel_values 关键字参数，需直接传入 tensor: model(pixel_values)。