WebSSL-MAE 1B (ViT-1B) 昇腾 NPU 适配

1. 简介

本文档记录 facebook/webssl-mae1b-full2b-224 模型在华为昇腾 NPU 上的适配与验证结果。该模型是基于 ViT-1B 架构的视觉自监督学习模型（Masked Autoencoder），参数量为 1.14B。

运行模式为原生 PyTorch 推理（非 vLLM 服务模式），通过 torch_npu 将模型加载到昇腾 NPU 上进行推理加速。

相关获取地址：

• 权重下载地址（HuggingFace）：https://huggingface.co/facebook/webssl-mae1b-full2b-224
• 权重下载地址（ModelScope）：https://modelscope.cn/models/facebook/webssl-mae1b-full2b-224
• 权重下载地址（AtomGit）：https://atomgit.com/models/webssl-mae1b-full2b-224

2. 验证环境

3. 推理脚本

3.1 单次推理（inference.py）

cd /opt/atomgit/webssl-mae1b-full2b-224

# 单张图片推理（NPU）
python3 inference.py --image path/to/image.jpg --device npu

# CPU 推理
python3 inference.py --image path/to/image.jpg --device cpu

3.2 精度评测（CPU vs NPU）

# 自动下载模型并进行 CPU ↔ NPU 精度对比
python3 accuracy_run.py 2>&1 | tee -a log.txt

脚本会分别用 CPU 和 NPU 推理同一组随机图片，计算归一化均方误差（NMSE）。

3.3 性能基准测试

# NPU 性能测试（含 CPU 单次参考）
python3 accuracy_run_perf.py 2>&1 | tee -a log.txt

# 或使用评测所有项目脚本
python3 check_accuracy_run_perf.py 2>&1 | tee -a log.txt

4. 适配要点

该模型为标准 HuggingFace ViTModel 架构，无需修改模型代码即可在昇腾 NPU 上运行：

1. 设备映射：调用 model.to("npu:0") 即可将模型参数和计算迁移到 NPU
2. 算子兼容性：ViT 的标准算子（nn.Linear, nn.LayerNorm, nn.Softmax, nn.GELU 等）在 torch_npu 上均有原生支持，无需替换
3. 同步控制：性能测试时通过 torch.npu.synchronize() 确保准确计时
4. 高速缓存：启用 expandable_segments:True 优化显存分配

环境变量推荐配置：

export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export TASK_QUEUE_ENABLE=1

5. 精度评测结果

测试条件：2 张随机 RGB 图片（224×224），CPU vs NPU last_hidden_state 逐元素对比。

NMSE 远低于 1% 阈值，精度完全满足要求。

6. 性能参考

测试条件：单张 224×224 图片，NPU 预热 5 次后采样 20 次取平均，CPU 单次计时。

7. 输出文件说明

适配完成后，所有关键文件位于 webssl-mae1b-full2b-224-ascend/ 目录：

8. 注意事项

• 首次运行需要下载模型权重（~4.24GB），可通过 HuggingFace、ModelScope 或 AtomGit 获取
• CPU 推理 1.14B 模型耗时较长（~25 秒/张），建议仅在精度校验时使用
• NPU 推理前需确保 torch_npu 正确安装且 NPU 驱动正常
• 当前适配方案适用于标准 ViTModel，无需 trust_remote_code