#NPU #昇腾 #Cross-Encoder #Reranker #BGE #文本重排序bge-reranker-base 是北京智源人工智能研究院(BAAI)开发的轻量级重排序模型,基于 BERT-base 架构,适用于:
| 项目 | 状态 |
|---|---|
| 推理引擎 | transformers + torch_npu |
| 运行设备 | Ascend 910 (NPU) |
| 精度验证 | ✅ 通过(NPU vs CPU 平均相对误差 0.14%) |
| 功能验证 | ✅ 通过 |
| 内存占用 | 约 1.04 GB |
| 组件 | 版本 |
|---|---|
| Python | 3.8+ |
| PyTorch | 2.9.0+ |
| torch_npu | 2.9.0.post1+ |
| transformers | 4.57.6+ |
| Ascend NPU 驱动 | 25.5.2+ |
# 安装 torch_npu(如果已有可跳过)
pip install torch torch_npu -i https://pypi.tuna.tsinghua.edu.cn/simple
# 安装 transformers
pip install transformers -i https://pypi.tuna.tsinghua.edu.cn/simplepython3 inference.py [选项]| 参数 | 默认值 | 说明 |
|---|---|---|
--model_path | BAAI/bge-reranker-base | 模型路径或 ModelScope/HuggingFace ID |
--device | 0 | NPU 设备 ID |
--max_length | 512 | 最大序列长度 |
--batch_size | 8 | 批处理大小 |
--num_runs | 10 | 性能测试运行次数 |
--warmup_runs | 3 | 预热运行次数 |
# 基本推理
python3 inference.py
# 使用本地模型
python3 inference.py --model_path /path/to/bge-reranker-base
# 性能基准测试(20 次运行)
python3 inference.py --num_runs 20 --warmup_runs 5使用 eval_accuracy.py 对同一模型同时在 CPU 和 NPU 上运行推理,对比输出结果的绝对误差和相对误差。
Ascend 910 NPU vs CPU 精度对比(10 组测试用例):
| # | Score(CPU) | Score(NPU) | 差值 | 相对误差 |
|---|---|---|---|---|
| 1 | 9.1247 | 9.1262 | 0.0016 | 0.02% |
| 2 | -10.1945 | -10.1945 | 0.0000 | 0.00% |
| 3 | 8.9284 | 8.9336 | 0.0052 | 0.06% |
| 4 | -1.0015 | -0.9915 | 0.0100 | 1.00% |
| 5 | 4.9992 | 5.0103 | 0.0111 | 0.22% |
| 6 | -10.1938 | -10.1938 | 0.0000 | 0.00% |
| 7 | 5.6669 | 5.6680 | 0.0011 | 0.02% |
| 8 | -10.1932 | -10.1932 | 0.0000 | 0.00% |
| 9 | 9.7237 | 9.7244 | 0.0006 | 0.01% |
| 10 | 6.6702 | 6.6772 | 0.0070 | 0.10% |
精度结论:AvgMaxRelErr 为 None%,MaxRelErr 为 1.00%,均低于 1% 精度要求,精度验证通过(PASS)。
精度统计:
✅ 精度验证通过!NPU 推理结果与 CPU 高度一致,排序结果完全一致。第 4 组用例误差 1.00% 为数值精度边界情况,不影响排序效果。
测试环境: Ascend 910 NPU | PyTorch 2.9.0 | torch_npu 2.9.0.post1 | batch_size=10
| 指标 | 数值 |
|---|---|
| 模型加载耗时 | 1.74s |
| 平均推理时间 | 8.02 ms/批 |
| 吞吐量 | 1247.64 序列/秒 |
| NPU 内存占用 | 1.04 GB |
| NPU 内存预留 | 1.10 GB |
import torch
import torch_npu
from transformers import AutoModelForSequenceClassification, AutoTokenizer
# 加载模型
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-base")
model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-base")
model.eval()
model = model.npu()
# 准备数据
pairs = [
("什么是机器学习?", "机器学习是人工智能的一个分支。"),
("什么是机器学习?", "今天天气很好。"),
]
inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors="pt", max_length=512)
inputs = {k: v.npu() for k, v in inputs.items()}
# 推理
with torch.no_grad():
outputs = model(**inputs)
scores = outputs.logits.squeeze(-1)
print(scores) # 正值表示相关,负值表示不相关| 模型 | 参数量 | 推理时间 (ms/批) | 吞吐量 (序列/秒) | 内存 (GB) |
|---|---|---|---|---|
| bge-reranker-base | 278M | 8.02 | 1247.64 | 1.04 |
| bge-reranker-v2-m3 | 567M | 14.24 | 702.07 | 2.12 |
bge-reranker-base-npu/
├── README.md # 本文件(部署文档)
├── inference.py # NPU 推理脚本
└── eval_accuracy.py # 精度验证脚本首次运行需要下载模型权重(约 1.04 GB),建议预先下载:
modelscope download --model BAAI/bge-reranker-base --local_dir /path/to/modelbge-reranker-base 仅需约 1.04 GB 显存,适合在大多数 NPU 环境运行。如仍有显存问题,请检查:
npu-smi info # 查看 NPU 显存使用情况基于现有评测数据,CPU 与 NPU 的 平均相对误差 精度误差为 0.14%,小于 1% 的精度要求。
本仓库提供完整的推理脚本,支持 CPU 和 NPU 双平台推理:
# NPU 推理
python3 inference.py --device npu
# CPU 推理
python3 inference.py --device cpu推理完成后会输出推理结果和耗时,表明模型在 NPU 上推理成功。
基于现有评测数据,CPU 与 NPU 的 平均相对误差 精度误差为 0.14%,小于 1% 的精度要求。