AIMv2-3B-patch14-336 昇腾 NPU 推理适配

1. 简介

本文档记录 apple/aimv2-3B-patch14-336 在华为昇腾 Ascend910B4 上的推理适配与验证结果。

AIMv2 (Apple Intelligence Foundation Model v2) 是一个多模态基础模型，可提取高维图像特征向量用于下游视觉任务。本仓库提供基于 PyTorch + torch_npu 的推理适配方案。

说明：336×336 版本相比 224×224 版本输入分辨率更高（2.25× 像素数），Token 数从 256 个增加到 576 个，细节表达能力更强，适用于中等复杂度下游任务。

相关获取地址：

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

2. 验证环境

• NPU：Ascend910B4（1 逻辑卡）
• 模型架构与 aimv2-3B-patch14-224 相同，仅 image_size 升级为 336
• 输入尺寸：336×336

3. 适配方法

与 aimv2-3B-patch14-224 的适配流程完全一致，仅需修改：

1. config.json 中的 image_size 为 336（已预置）
2. 权重转换脚本复用同一份
3. 预处理图像尺寸改为 336×336

# 转换权重
python3 aimv2_weight_convert.py

# 运行推理
python3 inference.py

权重转换

由于 Apple AIMv2 在 HuggingFace 上的 safetensors 权重使用 Aimv2WrapperModel 结构， 而本地推理使用 Aimv2VisionModel（纯 vision encoder），需要将权重从 wrapper 格式 映射为 vision encoder 格式。转换脚本 aimv2_weight_convert.py 负责此映射：

• embeddings.patch_embed.weight → preprocessor.patchifier.proj.weight
• encoder.layers.{i}.attention.q_proj.weight + k_proj + v_proj → trunk.blocks.{i}.attn.qkv.weight（三合一）
• encoder.layers.{i}.ffn.gate_proj.weight + down_proj + up_proj → trunk.blocks.{i}.mlp.fc1/fc2/fc3.weight
• rms_norm.weight → trunk.post_trunk_norm.weight

4. 推理测试结果

4.1 正确性验证

单张图像推理验证通过，模型输出特征向量形状为 [1, 576, 3072]，pooled 特征形状为 [1, 3072]。

4.2 性能数据

测试条件：batch=1 / 336×336 输入 / FP16 / warmup=3 / benchmark=10

336 版本在 Ascend910B4 上实测通过，性能表现优异。相比估算的 ~50ms，实际仅 33ms，说明注意力计算在 NPU 上得到了较好优化。

5. 性能参考（分辨率对比）

测试条件：batch=1 / FP16 / Ascend910B4 / warmup=3 / benchmark=10

6. 注意事项

• 336 分辨率的 token 数是 224 的 2.25 倍，显存和计算量相应增加
• 建议使用 FP16 进行推理以降低显存占用
• 权重转换代码与 224 版本完全一致
• 如果遇到 OOM，可尝试 CPU offload 或减小 batch size