VTP-Small-f16d64 on Ascend NPU

1. 简介

本文档记录 VTP-Small-f16d64 在华为昇腾 NPU 环境下的端到端适配、推理与精度验证结果。

VTP-Small-f16d64 是 MiniMax 开源的视觉 Tokenizer 预训练模型（Visual Tokenizer Pre-training），支持以下三种能力：

图像重建（Auto-encoder Reconstruction）：将输入图像编码为 latent 特征后解码重建为像素空间
CLIP 零样本分类（Zero-shot Classification）：图文对比学习特征提取与 zero-shot 推理
SSL 特征提取（Self-Supervised Feature Extraction）：自监督视觉特征提取，支持 intermediate layers 多层级特征

模型配置文件说明：vision_feature_bottleneck=64，即 ViT 特征经过 bottleneck 压缩到 64 维；decoder_embed_dim=384，解码器为 12 层 transformer 架构，最终通过 PixelShuffle 上采样 16 倍恢复图像。

原始仓库：https://github.com/MiniMax-AI/VTP

权重获取方式：

源	地址
HuggingFace 镜像	https://hf-mirror.com/MiniMax-VL/VTP-Small-f16d64
ModelScope	https://modelscope.cn/models/MiniMax-VL/VTP-Small-f16d64

2. 验证环境

组件	版本
`torch`	`2.9.0+cpu`
`torch-npu`	`2.9.0.post1+gitee7ba04`
`transformers`	`4.57.6`
`timm`	`1.0.27`
`omegaconf`	`2.3.0`
`torchmetrics`	`1.6.2`
`scipy`	`1.15.2`
`pytorch_fid`	`0.3.0`

NPU：1 逻辑卡（Ascend 910B4）
CANN 版本：8.5.1
模型路径：/opt/atomgit/VTP-Small-f16d64/weights
工作目录：/opt/atomgit/VTP-Small-f16d64/VTP-Small-f16d64ascend/

3. 依赖安装

推荐使用华为 PyPI 源加速安装：

pip install torch torchvision torchaudio timm omegaconf torchmetrics scipy pytorch_fid transformers \
  -i https://repo.huaweicloud.com/repository/pypi/simple

如果已存在 torch_npu 环境，仅需补充安装 VTP 仓库依赖：

pip install timm omegaconf torchmetrics scipy pytorch_fid \
  -i https://repo.huaweicloud.com/repository/pypi/simple

注意：bpe_simple_vocab_16e6.txt.gz 文件位于 tools/ 目录下，CLIP tokenizer 会自动搜索该路径。

4. 推理脚本

推理脚本 inference.py 支持以下三种任务的一键执行：

cd /opt/atomgit/VTP-Small-f16d64
python VTP-Small-f16d64ascend/inference.py

脚本自动检测可用的 NPU/CUDA/CPU 设备，依次执行：

图像重建 — 读取 weights/figures/dog.png，编码为 latents 再解码重建，保存到 output/reconstructed.png
CLIP zero-shot 分类 — 对输入图像提取特征，与 ["a diagram", "a dog", "a cat", "a person"] 文本特征计算相似度
SSL 特征提取 — 提取 Vision Encoder 最后一层 CLS / Patch tokens，以及 intermediate 4 层特征

运行命令如下：

export TORCH_HOME=/tmp/torch_home
mkdir -p $TORCH_HOME
python VTP-Small-f16d64ascend/inference.py

执行结果同时输出到终端和 VTP-Small-f16d64ascend/log.txt。

一键验证

提供 run_verify.sh 脚本，包含环境检查、推理、精度评测全流程：

bash VTP-Small-f16d64ascend/run_verify.sh

5. 精度评测

精度评测脚本 accuracy_eval.py 分别使用 NPU 和 CPU 两个设备运行相同的推理流程，逐任务对比输出张量的差异，计算相对误差百分比。

评测方法：

相对误差（Relative Error）：mean(|a - b| / |b|) * 100%，要求 < 1%
余弦相似度（Cosine Similarity）：衡量特征向量的方向一致性，越接近 1 越好
MSE / PSNR：针对重建图像的质量指标

python VTP-Small-f16d64ascend/accuracy_eval.py

脚本会自动将评测结果写入 VTP-Small-f16d64ascend/log.txt，并将汇总数据保存为 accuracy_summary.json。

典型精度结果

评测项	相对误差 (%)	余弦相似度	是否通过
图像重建 (Recon)	`< 0.01`	`> 0.9999`	✅
CLIP 图像特征	`< 0.01`	`> 0.9999`	✅
CLIP 文本特征	`< 0.01`	`> 0.9999`	✅
CLIP 分类概率	`< 0.01`	—	✅
SSL CLS Token	`< 0.01`	`> 0.9999`	✅
SSL Patch Tokens	`< 0.01`	`> 0.9999`	✅

以上结果为 NPU vs CPU 对比。由于 VTP-Small 是纯 PyTorch 视觉模型，NPU 上的 torch.autocast 与 CPU 上的 FP32 推理在数值上保持一致，误差远小于 1%。

6. 性能参考

测试条件：Ascend 910B4 (1 逻辑卡)，输入图像 256×256。

任务	耗时
模型加载	`2.36 s`
图像重建 (Enc + Dec)	`7.961 s`
CLIP zero-shot	`0.145 s`
SSL 特征提取	`0.060 s`

性能数据来源于单次运行，实际性能会受 NPU 负载、内存带宽等因素影响。

7. 适配要点

7.1 核心改动

文件	说明
`inference.py`	昇腾 NPU 推理入口，支持三种任务
`accuracy_eval.py`	精度评测脚本，NPU vs CPU 逐任务对比
`utils_ascend.py`	Ascend 辅助工具函数
`run_verify.sh`	一键验证脚本

7.2 亲和算子替换

通过 named_replace 机制在模型加载前后替换以下算子，提升 NPU 推理性能：

nn.GELU → torch.nn.GELU（使用 approximate="tanh"）
nn.LayerNorm → AscendLayerNorm（torch_npu 融合算子）

详细替换记录见 fusion_result.json。

8. 注意事项

PixelShuffle 警告：推理过程中可能出现 as_strided 相关 UserWarning，这是 torch_npu 对 nn.PixelShuffle 的内部提示，不影响重建精度与功能正确性。如遇到精度问题，可尝试设置 torch.npu.config.allow_internal_format = False。
Tokenizer BPE 文件：inference.py 依赖 bpe_simple_vocab_16e6.txt.gz，请确保该文件位于 tools/ 目录下，或在 vtp/tokenizers/ 的搜索路径内。
ACL Graph 不适用：VTP 是纯 PyTorch 视觉模型，不经过 vLLM / ACL Graph 链路，因此无需配置 cudagraph_capture_sizes 等参数。
多卡扩展：当前脚本为单 NPU 推理。如需多卡并行，可参考 torch.nn.DataParallel 或 torch.distributed 自行扩展。
权重下载限制：如果 GitHub/HuggingFace 网络不通，请使用 ModelScope 镜像或华为云 OBS 下载模型权重。