AIMv2-1B-Patch14-336 on Ascend NPU

模型简介

aimv2-1B-patch14-336 是 Apple 发布的 AIMv2 系列视觉 Transformer (ViT) 模型之一，输入分辨率为 336×336，patch size 为 14×14。隐藏层维度 2048，16 头注意力，参数量 1.24B。

本仓库提供了该模型在 华为昇腾 Ascend910B4 NPU 上的推理适配方案，基于原生 transformers 4.57.6 实现，无需 trust_remote_code。

模型架构

• Patch Embed: Conv2d (patch_size=14) + RMSNorm + 可学习位置编码
• Transformer Encoder: 24 层, hidden_size=2048, 16 头注意力, SwiGLU FFN (intermediate_size=5632)
• 输出: last_hidden_state [1, 576, 2048]

相关地址

• 权重下载（ModelScope）：https://modelscope.cn/models/apple/aimv2-1B-patch14-336
• 权重下载（HuggingFace）：https://huggingface.co/apple/aimv2-1B-patch14-336

验证环境

• 模型路径：/opt/atomgit/.cache/modelscope/hub/models/apple/aimv2-1B-patch14-336
• 推理设备：npu:0

模型加载

from transformers import AutoImageProcessor
from transformers.models.aimv2.modeling_aimv2 import Aimv2VisionModel

model_path = "/path/to/aimv2-1B-patch14-336"
processor = AutoImageProcessor.from_pretrained(model_path, trust_remote_code=True)
model = Aimv2VisionModel.from_pretrained(model_path)
model.eval()
model = model.to("npu:0")

说明：ModelScope 提供的自定义 modeling_aimv2.py 使用合并 QKV 投影，与权重文件（分离 Q/K/V 投影）不兼容。transformers 4.57.6 原生 Aimv2VisionModel 使用分离投影，可直接加载权重。

推理验证

推理脚本

# 随机图片推理
python3 inference.py --model 336 --device npu:0

# 真实图片推理
python3 inference.py --model 336 --device npu:0 --image /path/to/image.jpg

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

验证结果

[INFO] Loading model from: /path/to/aimv2-1B-patch14-336
[INFO] Parameters: 1235.61M
[INFO] Device: npu:0
[INFO] Input shape: [1, 3, 336, 336]
[INFO] Inference time: 37.53 ms
[INFO] Output shape: [1, 576, 2048]
[DONE] Inference completed successfully

精度评测

评测方法：CPU (FP32) 参考 vs NPU (FP32) 推理，逐元素对比。

NPU 与 CPU 输出余弦相似度 > 0.9999，平均相对误差 < 0.3%，满足精度要求。

性能参考

测试条件：Ascend910B4 × 1, FP32, batch size 1~8, 20 轮取平均。

评测材料

• 精度评测：evaluate.py + eval_results/accuracy.json
• 性能基准：evaluate.py + eval_results/performance.json
• 评测日志：eval_results/run.log

注意事项

1. 自定义代码兼容性：ModelScope 下载的 modeling_aimv2.py 与 safetensors 权重存在结构不匹配（合并 QKV vs 分离 QKV），请使用原生 Aimv2VisionModel 加载。
2. NPU 浮点差异：NPU 的浮点累加顺序与 CPU 不同，会导致元素级微小差异（一般 < 0.3%），但不影响下游任务精度。
3. use_head 参数：当前权重 use_head=false，模型仅输出 last_hidden_state，pooler_output 为 None。