wd-swinv2-tagger-v3

1. 简介

本文档记录 fireicewolf/wd-swinv2-tagger-v3 (Swin 图像分类模型) 在昇腾 NPU（Ascend910）环境的快速部署与验证结果。

SwinV2-Base 图像分类模型，基于 PyTorch + timm 框架，支持一键加载推理。

2. 验证环境

组件	版本
`torch`	`2.5.1`
`torch_npu`	`2.5.1`
`timm`	`>=4.48.0`
`CANN`	`8.5.RC1`

NPU：Ascend910（单卡）
输入尺寸：?
类别数：?
推理框架：PyTorch + timm

3. 快速部署

3.1 环境准备

pip install timm safetensors torch torchvision pillow

3.2 推理代码

import torch
import timm
from safetensors.torch import load_file
from PIL import Image
from torchvision import transforms

device = "npu:0" if torch.npu.is_available() else "cpu"
model = timm.create_model("swinv2_base_window8_256", pretrained=False, num_classes=10861, img_size=448, window_size=14)
state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict, strict=True)
model = model.to(device).eval()

transform = transforms.Compose([
    transforms.Resize(448),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
img = Image.open("image.jpg")
t = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
    output = model(t)
pred = output.argmax(-1).item()
print(f"Predicted class: {pred}")

4. Smoke 验证

python3 inference.py

验证结果：

模型成功加载到 npu:0
输出预测类别（整数索引）
推理过程无报错

5. 性能参考

测试条件：FP32 / batch=1 / warmup=5 / timed=50 runs，Ascend910 单卡。

指标	数值
平均推理时间	`23.65 ms`
测试次数	`50`

6. 精度评测

NPU 与 CPU 输出对比，使用 4 张纯色测试图（RGB 255/0、0/255/0、0/0/255、128/128/128），比较 logits 一致性。

指标	数值
Top-1 一致性	`4/4`
Top-5 一致性	`4/4`
最大 logits 相对误差	`0.453` %
平均 KL 散度	`6.8e-05`
结论	`PASS`

7. 注意事项

AutoImageProcessor 会自动匹配模型配置中的 image_size，无需手动设置
Swin Transformer 使用 SwinForImageClassification，AutoModelForImageClassification 可自动识别
NPU 推理结果与 CPU 在 Top-1/Top-5 完全一致，最大 logits 相对误差 < 1%（判定阈值）
22k 预训练模型（ImageNet-22K）id2label 为 21841 类，使用 AutoImageProcessor 自动处理输入分辨率

组件

版本

torch

2.5.1

torch_npu

2.5.1

timm

>=4.48.0

CANN

8.5.RC1

3. 快速部署

3.1 环境准备

pip install timm safetensors torch torchvision pillow

3.2 推理代码

import torch
import timm
from safetensors.torch import load_file
from PIL import Image
from torchvision import transforms

device = "npu:0" if torch.npu.is_available() else "cpu"
model = timm.create_model("swinv2_base_window8_256", pretrained=False, num_classes=10861, img_size=448, window_size=14)
state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict, strict=True)
model = model.to(device).eval()

transform = transforms.Compose([
    transforms.Resize(448),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
img = Image.open("image.jpg")
t = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
    output = model(t)
pred = output.argmax(-1).item()
print(f"Predicted class: {pred}")

指标

数值

平均推理时间

23.65 ms

测试次数

50

指标

数值

Top-1 一致性

4/4

Top-5 一致性

4/4

最大 logits 相对误差

0.453 %

平均 KL 散度

6.8e-05

结论

PASS

7. 注意事项

AutoImageProcessor 会自动匹配模型配置中的 image_size，无需手动设置

Swin Transformer 使用 SwinForImageClassification，AutoModelForImageClassification 可自动识别

NPU 推理结果与 CPU 在 Top-1/Top-5 完全一致，最大 logits 相对误差 < 1%（判定阈值）

22k 预训练模型（ImageNet-22K）id2label 为 21841 类，使用 AutoImageProcessor 自动处理输入分辨率