ConvNeXt Base (CLIP LAION-2B) on Ascend NPU

1. 简介

本文档记录 convnext_base.clip_laion2b 在华为昇腾 Ascend910B4 NPU 上的适配部署与验证结果。

convnext_base.clip_laion2b 是一个 88M 参数的 ConvNeXt 架构视觉模型，基于 CLIP 在 LAION-2B 数据集上训练，输出 640 维多模态嵌入向量。模型使用标准 PyTorch 算子（Conv2d、LayerNorm、GELU、Linear），无需任何算子替换即可在 Ascend NPU 上原生运行。

2. 验证环境

组件	版本
`PyTorch`	`2.9.0+cpu`
`torch_npu`	`2.9.0.post1+gitee7ba04`
`timm`	`1.0.27`
`CANN`	`8.5.1`

NPU：Ascend910B4 x 1 逻辑卡
模型路径：/opt/atomgit/convnext_base.clip_laion2b_modelscope/timm/convnext_base.clip_laion2b/model.safetensors
输入尺寸：(3, 224, 224)
输出维度：640

3. 快速部署

3.1 环境准备

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0

pip install timm safetensors pillow numpy

注意：ASCEND_RT_VISIBLE_DEVICES 必须设为 0（NPU 本地索引），而非物理设备号。

3.2 推理命令

# NPU 推理
python3 inference.py --device npu --image input.jpg

# CPU 推理（对比用）
python3 inference.py --device cpu --image input.jpg

3.3 代码示例

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu

from timm.models.convnext import convnext_base
from safetensors.torch import load_file
from PIL import Image
from timm.data import resolve_data_config, create_transform

# 加载模型
model = convnext_base(pretrained=False, num_classes=640)
state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict, strict=False)
model = model.to("npu")
model.eval()

# 预处理
cfg = resolve_data_config({}, model=convnext_base(pretrained=False))
transform = create_transform(**cfg)
img = transform(Image.open("input.jpg").convert("RGB")).unsqueeze(0)

# 推理
with torch.no_grad():
    output = model(img.to("npu"))
print(output.shape)  # [1, 640]
print(output[0].topk(5).indices)  # Top-5 分类

4. Smoke 验证

# 验证 NPU 可用性
python3 -c "import torch, torch_npu; print(torch.npu.is_available(), torch.npu.get_device_name(0))"

# 运行推理
python3 inference.py --device npu --image test_input.jpg

预期输出：

Device: npu
NPU available: True, Device: Ascend910B4
Input config: {'input_size': (3, 224, 224), ...}
Image: test_input.jpg -> (1, 3, 224, 224)
Avg inference time: 50.62 ms
Output shape: (1, 640)
Top-5 indices: [493, 502, 528, 320, 505]

5. 性能参考

测试条件：batch=1，随机输入，warmup=3，runs=20。

指标	CPU	NPU
单张延迟 (bs=1)	`1747.73 ms`	`18.59 ms`
吞吐量 (bs=1)	`0.57 img/s`	`53.78 img/s`
加速比 (bs=1)	`1x`	~94x

NPU 多 Batch 性能

Batch Size	延迟 (ms)	吞吐量 (img/s)
1	`18.59`	`53.78`
2	`18.51`	`108.05`
4	`19.54`	`204.69`
8	`24.17`	`331.02`
16	`35.23`	`454.22`
32	`61.18`	`523.01`

6. 精度评测

使用 50 组随机输入对比 CPU (float32) 与 NPU (float32) 输出差异。

指标	数值
测试样本数	`50`
Max 绝对差值	`1.12e-02` (< 1.2%)
Mean 绝对差值	`1.27e-03`
Cosine Similarity (avg)	`0.99998967`
Cosine Similarity (min)	`0.99998684`
Shape 一致性	`完全一致`
精度判定	通过（Cosine Similarity > 0.9999）

注：Ascend910B4 不支持 fp64，torch_npu 自动将 DoubleTensor 降级为 FloatTensor。微小差异来自 fp32 精度下不同硬件后端（CPU vs NPU）的浮点运算顺序差异，属正常现象。

7. 注意事项

ASCEND_RT_VISIBLE_DEVICES 配置：必须设为 NPU 本地索引 0，而非物理设备号。本环境 NPU 物理 ID 为 7，但运行时只识别 1 个设备，索引为 0，设 7 会导致 aclrtGetDeviceCount 返回错误。
环境变量必须提前 source：source /usr/local/Ascend/ascend-toolkit/set_env.sh 必须在 Python 启动前执行，否则 torch_npu 无法完成 aclInit。
transfer_to_npu 日志：首次导入 from torch_npu.contrib import transfer_to_npu 时会打印 CUDA API 替换列表，此为正常行为，不影响功能。
fp64 降级：Ascend910 系列不支持 double 精度，若有代码使用 torch.DoubleTensor 会静默降级为 torch.FloatTensor，不影响推理精度。
日志目录权限：若 /home/atomgit/ascend/log 不可写，会出现 LOG_WARNING，可忽略，不影响推理。