ViT (Vision Transformer) 昇腾 NPU 迁移文档

import timm
model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=1000)

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export HF_ENDPOINT=https://hf-mirror.com

python3 -c "
import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
import timm, torch
model = timm.create_model('vit_base_patch16_224', pretrained=True)
model.eval()
torch.save(model.state_dict(), '/home/zy/vit-a3/vit_base_patch16_224_weights.pth')
print('Weights saved. Params:', sum(p.numel() for p in model.parameters()) / 1e6, 'M')
"

mkdir -p /home/zy/vit-a3/imagenet_val
cd /home/zy/vit-a3/imagenet_val

for i in $(seq 0 14); do
  fname=$(printf "validation-%05d-of-00015.parquet" $i)
  wget -O "$fname" \
    "https://hf-mirror.com/datasets/Tsomaros/Imagenet-1k_validation/resolve/main/data/$fname"
done

python3 -c "
import pyarrow.parquet as pq, os
from collections import Counter
total, labels = 0, []
for f in sorted(os.listdir('/home/zy/vit-a3/imagenet_val')):
    if f.startswith('validation-') and f.endswith('.parquet'):
        t = pq.read_table(os.path.join('/home/zy/vit-a3/imagenet_val', f))
        total += t.num_rows
        labels.extend(t.column('label').to_pylist())
lc = Counter(labels)
print(f'Total: {total}, Classes: {len(lc)}, Per class: {Counter(lc.values())}')
"

docker run -it -u root -d --net=host \
--privileged \
--ipc=host \
--device=/dev/davinci_manager \
--device=/dev/devmm_svm \
--device=/dev/hisi_hdc \
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/sbin:/usr/local/sbin \
-v /usr/local/Ascend/driver/tools/hccn_tool:/usr/local/Ascend/driver/tools/hccn_tool \
--name zy-a3-ubuntu \
quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11 \
/bin/bash

# 激活 CANN 环境（每次执行脚本前必须执行）
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 安装 timm（模型实现）
pip install --no-deps timm==1.0.19 -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装 torchvision（图像预处理，匹配 torch 2.1.0）
pip install --no-deps torchvision==0.16.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装 transformers 及其依赖
pip install --no-deps "transformers>=4.44.0" -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install --no-deps "tokenizers>=0.19,<0.20" -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装数据处理依赖
pip install --no-deps Pillow pyarrow numpy -i https://pypi.tuna.tsinghua.edu.cn/simple

# 检查 Python 版本
python3 --version

# 检查 NPU 状态
npu-smi info

# 检查驱动版本
cat /usr/local/Ascend/driver/version.info

# 激活 CANN 环境
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 验证 NPU 可用
python3 -c "
import torch, torch_npu
print('torch:', torch.__version__)
print('torch_npu:', torch_npu.__version__)
print('NPU available:', torch.npu.is_available())
print('NPU count:', torch.npu.device_count())
"

torch: 2.1.0
torch_npu: 2.1.0.post17
NPU available: True
NPU count: 16

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

pip install --no-deps timm==1.0.19 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install --no-deps torchvision==0.16.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install --no-deps Pillow pyarrow numpy -i https://pypi.tuna.tsinghua.edu.cn/simple

python3 -c "
import timm, torchvision, PIL, pyarrow
print('timm:', timm.__version__)
print('torchvision:', torchvision.__version__)
print('Pillow:', PIL.__version__)
print('pyarrow:', pyarrow.__version__)
print('All imports OK')
"

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export HF_ENDPOINT=https://hf-mirror.com

python3 -c "
import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
import timm, torch
model = timm.create_model('vit_base_patch16_224', pretrained=True)
model.eval()
torch.save(model.state_dict(), '/home/zy/vit-a3/vit_base_patch16_224_weights.pth')
print('Weights saved. Params:', sum(p.numel() for p in model.parameters()) / 1e6, 'M')
"

Weights saved. Params: 86.567656 M

mkdir -p /home/zy/vit-a3/imagenet_val
cd /home/zy/vit-a3/imagenet_val

for i in $(seq 0 14); do
  fname=$(printf "validation-%05d-of-00015.parquet" $i)
  echo "Downloading $fname ..."
  wget -q --show-progress -O "$fname" \
    "https://hf-mirror.com/datasets/Tsomaros/Imagenet-1k_validation/resolve/main/data/$fname"
done

echo "Download complete"
ls -lh *.parquet | wc -l
du -sh .

python3 -c "
import pyarrow.parquet as pq, os
from collections import Counter
total, labels = 0, []
for f in sorted(os.listdir('/home/zy/vit-a3/imagenet_val')):
    if f.startswith('validation-') and f.endswith('.parquet'):
        t = pq.read_table(os.path.join('/home/zy/vit-a3/imagenet_val', f))
        total += t.num_rows
        labels.extend(t.column('label').to_pylist())
lc = Counter(labels)
print(f'Total: {total}, Classes: {len(lc)}, Per class: {Counter(lc.values())}')
"

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export OPENBLAS_NUM_THREADS=1
export OMP_NUM_THREADS=1

python3 /home/zy/vit-a3/benchmark_npu.py \
    --data-dir /home/zy/vit-a3/imagenet_val \
    --model vit_base_patch16_224 \
    --weights /home/zy/vit-a3/vit_base_patch16_224_weights.pth \
    --batch-size 64 \
    --device-id 0 \
    --warmup 20 \
    --output /home/zy/vit-a3/benchmark_npu_results.json

NPU available: True
NPU count: 16
Using NPU: 0 - Ascend910_9362
Loading model: vit_base_patch16_224
Model loaded. Parameters: 86.57M
Loading ImageNet validation set from /home/zy/vit-a3/imagenet_val
Loaded 50000 images, 1000 classes
Running NPU benchmark (batch_size=64, warmup=20)...
  Batch 50/782,  Top-1: 87.81%, Top-5: 98.22%
  Batch 100/782, Top-1: 88.89%, Top-5: 98.19%
  ...
  Batch 782/782, Top-1: 85.22%, Top-5: 97.47%

=== NPU Benchmark Results ===
Device: Ascend910_9362
Top-1 Accuracy: 85.22%
Top-5 Accuracy: 97.47%
Avg Latency: 0.81 ms/sample
Throughput: 1232.63 samples/s
Total Time: 225.67 s
Results saved to /home/zy/vit-a3/benchmark_npu_results.json

cat /home/zy/vit-a3/benchmark_npu_results.json

{
  "top1_accuracy": 85.218,
  "top5_accuracy": 97.472,
  "avg_latency_ms": 0.811,
  "throughput_fps": 1232.63,
  "total_samples": 50000,
  "warmup_batches": 20,
  "device": "NPU:0",
  "npu_name": "Ascend910_9362",
  "total_time_s": 225.67
}

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export OPENBLAS_NUM_THREADS=1
export OMP_NUM_THREADS=1

python3 /home/zy/vit-a3/inference_demo_npu.py \
    --num-samples 20 \
    --stride 50 \
    --top-k 5 \
    --output /home/zy/vit-a3/inference_demo_results.json

--- Image 1/20 [CORRECT] ---
  Ground Truth: [0] tench
  Prediction:   [0] tench  (conf: 53.03%)
  Top-5 predictions:
    1. [   0] tench                                     53.03%
    2. [ 394] sturgeon                                  10.28%
    ...

--- Image 3/20 [CORRECT] ---
  Ground Truth: [2] great white shark
  Prediction:   [2] great white shark  (conf: 74.37%)
  ...

Summary: 20/20 correct (100.0%)
Detailed results saved to: inference_demo_results.json

import torch_npu                           # 注册 NPU 后端
device = torch.device(f'npu:{device_id}')  # 选择 NPU 设备
model = model.to(device)                   # 模型迁移到 NPU
torch.npu.synchronize()                    # NPU 同步（精确计时）