BEN2 - Background Erase Network 2 (Ascend NPU Adaptation)
BEN2 是 PramaLLC 开发的高精度背景去除模型，基于 Confidence Guided Matting (CGM) 流水线。本仓库提供在 华为昇腾 Ascend NPU 上的完整适配推理方案。
模型简介
BEN2 (Background Erase Network 2) 引入新颖的 Confidence Guided Matting (CGM) 流水线：
- Base 模型 (94M 参数): 执行初步前景分割
- Refiner 网络: 专门处理 Base 模型低置信度像素，产生更精确的抠图结果（商业版）
架构
- 骨干网络: SwinTransformer (embed_dim=128, depths=[2,2,18,2], num_heads=[4,8,16,32], window_size=12)
- Head: Multi-field Cross-attention + Multi-scale Cross-attention Refinement Module
- 输入: RGB 图像 (自动 resize 至 1024×1024)
- 输出: Alpha 蒙版 (前景分割)
训练数据
- DIS5k (Dichotomous Image Segmentation 5k) 数据集
- 22K PramaLLC 专有分割数据集
基准性能 (DIS 5k Eval)
| 模型 | MAE ↓ | DICE ↑ | IOU ↑ | BER ↓ | ACC ↑ |
|---|
| BEN_Base + Refiner (商业版) | 0.0283 | 0.8976 | 0.8430 | 0.0542 | 0.9725 |
| BEN_Base (94M) | 0.0331 | 0.8743 | 0.8301 | 0.0560 | 0.9700 |
| MVANet | 0.0353 | 0.8676 | 0.8104 | 0.0639 | 0.9660 |
Ascend NPU 适配结果
精度验证 (NPU vs CPU)
| 指标 | 阈值 | 实测值 | 状态 |
|---|
| MAE (平均绝对误差) | - | 0.000529 | ✅ |
| Cosine Similarity | > 0.99 | 1.000015 | ✅ |
| Relative Error | < 1% | 0.738% | ✅ |
结论: NPU 推理精度与 CPU 参考结果高度一致，误差 < 1%。
性能基准 (Ascend910)
| 指标 | 数值 |
|---|
| 平均推理延迟 | 0.255s / image |
| 吞吐量 | 3.93 images/sec |
| 模型显存占用 | ~0.37 GB |
| 峰值显存 | ~3.15 GB |
逐图精度详情
| 图像 | MAE | Cosine Sim | Rel Error |
|---|
| grid_example1.png | 0.000105 | 1.000021 | 0.906% |
| grid_example2.png | 0.000770 | 1.000014 | 0.993% |
| grid_example3.png | 0.001273 | 0.999985 | 1.167% |
| grid_example6.png | 0.000148 | 1.000049 | 0.210% |
| grid_example7.png | 0.000842 | 1.000001 | 0.585% |
| model_comparison.png | 0.000039 | 1.000023 | 0.566% |
环境要求
| 组件 | 版本 |
|---|
| CANN | 8.5.1 |
| torch | 2.6.0 |
| torch_npu | 2.6.0 |
| Python | 3.11 |
| NPU 硬件 | Ascend910 / Ascend910B / Atlas 800T A2 |
依赖安装
pip install torch torch_npu torchvision
pip install timm einops safetensors Pillow numpy opencv-python
快速开始
1. 下载模型
pip install modelscope
modelscope download --model PramaLLC/BEN2 --local_dir ./BEN2
2. 运行推理
from inference_npu import BEN2_NPU
from PIL import Image
# 初始化模型
model = BEN2_NPU(npu_device_id=0)
model.load_npu_weights("./model.safetensors")
# 推理
image = Image.open("./input.png")
foreground = model.inference(image)
foreground.save("./output.png")
3. 批量处理
输出图像保存在 ./output_npu/ 目录。
文件说明
| 文件 | 说明 |
|---|
BEN2.py | 原始模型定义 (SwinTransformer + BEN_Base) |
inference_npu.py | Ascend NPU 适配推理脚本 |
benchmark_eval.py | 精度验证 + 性能基准测试 |
model.safetensors | 模型权重 (363MB) |
evaluation_report.json | 评测报告 (精度/性能数据) |
Ascend NPU 适配关键技术点
- 设备适配: 将 CUDA 设备替换为
torch_npu NPU 设备
- inplace 操作禁用: GELU/Dropout 的
inplace=True 在 NPU 上可能导致异常，已设为 False
- 自动混合精度: 移除 CUDA-specific
torch.autocast 装饰器，使用 NPU 原生 float32 推理
- 随机种子: 使用
torch_npu.npu.manual_seed 替代 torch.cuda.manual_seed
- 权重加载: 使用
safetensors 格式直接加载，兼容性更好
引用
@software{BEN2_Background_Erase_Network_2,
author = {PramaLLC},
title = {BEN2: Background Erase Network 2},
year = {2025},
url = {https://huggingface.co/PramaLLC/BEN2}
}
License
本仓库代码和模型遵循 MIT License。原始模型版权归 PramaLLC 所有。