冬

vit-age-classifier Ascend NPU 部署指南

项目简介

vit-age-classifier 是基于 Google ViT-Base-Patch16-224 的年龄分类模型，能够对输入的人脸图像进行年龄范围分类。模型将图像分类为 9 个年龄组：0-2、3-9、10-19、20-29、30-39、40-49、50-59、60-69、70+。

特性

支持 Ascend NPU 推理加速
CPU 与 NPU 精度对比测试（误差 < 1%）
9 类年龄分类
兼容 HuggingFace transformers
121.20 倍加速比

环境要求

硬件：华为 Ascend 910 系列 NPU
CANN：8.0.RC1 或更高版本
PyTorch：2.0+ 并带有 torch_npu
transformers：4.5+
Pillow、numpy

目录结构

vit-age-classifier-ascend/
├── inference.py          # 推理测试脚本
├── log.txt               # 测试日志
├── README.md             # 本文档
├── test_image.png        # 测试图像
├── inference_result.json # 推理结果
└── precision_result.json  # 精度测试结果

部署步骤

1. 进入容器

docker exec -it test-modelagent bash

2. 设置环境变量

source /usr/local/Ascend/ascend-toolkit/set_env.sh

3. 准备模型文件

模型文件位于 /data/ysws/agentsp/5-16/vit-age-classifier/ 目录下：

model.safetensors - 模型权重 (约 343MB)
pytorch_model.bin - PyTorch 权重备份
config.json - 模型配置
preprocessor_config.json - 图像预处理配置

4. 安装依赖

pip install transformers torch_npu pillow numpy -i https://pypi.huaweicloud.com/repository/pypi/simple/

Usage

Method 1: Normal Inference Mode

Run the inference script for age classification:

cd /data/ysws/agentsp/5-16/vit-age-classifier-ascend/

python3 inference.py

python3 inference.py --mode inference

Method 2: Accuracy Test Mode (CPU vs NPU)

Run the accuracy comparison test:

cd /data/ysws/agentsp/5-16/vit-age-classifier-ascend/

python3 inference.py --mode precision_test

命令行参数说明

参数	说明	默认值
`--mode`	测试模式: all, inference 或 precision_test	`all`

测试验证

精度测试结果

指标	实测值	阈值	状态
最大相对误差	0.2793%	< 1.00%	PASS
CPU 推理时间	1.683s	-	-
NPU 推理时间	0.014s	-	-
加速比	121.20x	> 1x	PASS

推理结果示例

输入: 224x224 RGB 人脸图像

输出:

类别 ID: 3
年龄范围: 20-29
Logits 维度: 9

年龄组	ID
0-2	0
3-9	1
10-19	2
20-29	3
30-39	4
40-49	5
50-59	6
60-69	7
70+	8

测试日志

ViT Age Classifier NPU Test
Model: google/vit-base-patch16-224-in21k (age classification)
Output: /data/ysws/agentsp/5-16/vit-age-classifier-ascend

============================================================
Inference Test (NPU)
============================================================
Device: npu:0
Loading model and processor...
Model loaded successfully
Input shape: torch.Size([1, 3, 224, 224])
Inference time: 5.470s
Predicted class: 3 (20-29)
Logits shape: torch.Size([1, 9])

============================================================
Precision Test (CPU vs NPU)
============================================================
NPU Device: npu:0
Loading model...
Input shape: torch.Size([1, 3, 224, 224])
Running on CPU...
Running on NPU...
CPU inference time: 1.683s
NPU inference time: 0.014s
Speedup: 121.20x
Max absolute error: 3.722869e-03
Max relative error: 0.2793% (threshold: 1.0%)
Status: PASS

============================================================
Precision Test Result: PASS
============================================================

============================================================
Test Complete!
============================================================

Python API 使用示例

基本推理

import torch
from PIL import Image
from transformers import ViTForImageClassification, AutoImageProcessor

MODEL_DIR = "/data/ysws/agentsp/5-16/vit-age-classifier"

processor = AutoImageProcessor.from_pretrained(MODEL_DIR)
model = ViTForImageClassification.from_pretrained(MODEL_DIR)
model = model.to("npu:0").eval()

image = Image.open("face.jpg").convert("RGB")
inputs = processor(images=image, return_tensors="pt")
inputs = {k: v.to("npu:0") for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)

logits = outputs.logits
predicted_class = logits.argmax(-1).item()
predicted_label = model.config.id2label[str(predicted_class)]
print(f"Predicted: {predicted_label}")

自定义图像处理

import numpy as np

array = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
image = Image.fromarray(array)
inputs = processor(images=image, return_tensors="pt")

模型结构

架构类型: ViT (视觉Transformer)
骨干网络: vit-base-patch16-224-in21k
编码器: 12层Transformer
隐藏层维度: 768
注意力头数: 12
参数量: ~86M
图像尺寸: 224x224

组件	说明
embeddings	补丁嵌入 + CLS标记
encoder	12层Transformer编码器
layernorm	层归一化
classifier	线性分类头 (768 -> 9)

推理参数配置

从 config.json 提取的关键参数:

{
  "hidden_size": 768,
  "intermediate_size": 3072,
  "num_attention_heads": 12,
  "num_hidden_layers": 12,
  "patch_size": 16,
  "image_size": 224,
  "num_channels": 3
}

常见问题

Q: 精度测试失败?

A: 检查 NPU 驱动是否正确安装。ViT 模型在 CPU 和 NPU 上的数值误差极小（< 0.3%），远低于 1% 阈值。

Q: 如何提高推理速度?

A: 首次推理会有编译开销。NPU 相比 CPU 有显著加速（121x），适合批量处理场景。

Q: 支持哪些图像格式?

A: 支持 PIL、numpy array、torch tensor 等常见格式。输入图像会自动 resize 到 224x224 并 normalize。

Q: 如何处理批量图像?

A: 修改 inputs 格式即可:

images = [Image.open(f) for f in image_files]
inputs = processor(images=images, return_tensors="pt")

参考链接

原始模型: https://huggingface.co/google/vit-base-patch16-224-in21k
ViT 论文: https://arxiv.org/abs/2010.11929
HuggingFace Transformers: https://huggingface.co/transformers

许可证

本项目遵循 Apache-2.0 许可证