all-MiniLM-L6-v2 on Ascend NPU

1. 简介

本文档记录 all-MiniLM-L6-v2（句子嵌入模型）在昇腾 NPU（Ascend 910B3）环境的适配部署与精度验证结果。

all-MiniLM-L6-v2 是一种轻量级句子嵌入模型，参数量约 22.7M，基于 MiniLM 架构（6 层 BERT），将输入句子映射到 384 维向量空间，适用于语义搜索、聚类和文本分类等下游任务。本项目完成该模型在昇腾 NPU 上的推理适配，并验证 NPU 与 CPU 推理结果的精度误差 < 1%。

相关地址：

• 权重下载地址（HuggingFace）：https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2
• 适配代码仓库：https://gitcode.com/dingdang666/all-MiniLM-L6-v2-NPU

2. 验证环境

3. 模型信息

4. Conda 环境安装

创建独立的 conda 环境并安装依赖（使用华为云 PyPI 镜像加速）：

# 创建 conda 环境
conda create -n minilm python=3.11 -y

# 激活环境
conda activate minilm

# 安装 PyTorch 及相关依赖（华为云镜像）
pip install torch==2.10.0 torchvision==0.25.0 --index-url https://repo.huaweicloud.com/repository/pypi/simple/

# 安装 torch_npu（根据环境选择对应版本）
pip install torch_npu==2.10.0 --index-url https://repo.huaweicloud.com/repository/pypi/simple/

# 安装 transformers 及其他依赖
pip install transformers safetensors --index-url https://repo.huaweicloud.com/repository/pypi/simple/

如果 HuggingFace 网络不通，设置镜像：

export HF_ENDPOINT=https://hf-mirror.com/

5. 推理执行

句子嵌入提取

# NPU 推理（默认）
python3 inference.py \
    --model_path /path/to/all-MiniLM-L6-v2 \
    --text "Your sentence here"

# CPU 推理
python3 inference.py \
    --model_path /path/to/all-MiniLM-L6-v2 \
    --text "Your sentence here" \
    --device cpu

精度与性能评测

python3 benchmark.py \
    --model_path /path/to/all-MiniLM-L6-v2 \
    --npu_device npu:0

评测结果日志将输出到 log.txt。

6. 参数说明

inference.py 参数

benchmark.py 参数

7. 精度评测

评测方法

使用同一段输入文本分别在 CPU（FP32）和 NPU（FP32）上运行推理，对比输出 CLS token 嵌入向量的差异。

评价指标说明

精度结果

判定结论

8. 性能数据

9. 注意事项

1. 首次加载：首次加载模型权重会从本地读取 safetensors 文件，耗时约 1 秒。
2. NPU 预热：NPU 首次推理包含编译优化，通常需要 1-2 轮预热才能达到稳定性能。脚本默认开启 3 轮预热。
3. 精度：本模型使用 FP32 精度，NPU 与 CPU 输出误差极低（~0.1%），不影响下游语义检索效果。
4. 权重文件：模型权重（.safetensors）不包含在适配仓库中，需单独下载。
5. 输出说明：脚本输出 last_hidden_state[:, 0]（CLS token）作为句子嵌入向量，维度为 384。
6. 轻量推理：22.7M 参数，NPU 推理仅需 ~0.01s，适合高吞吐场景。