mobilenetv2_100.ra_in1k 昇腾 NPU 适配

pip install torch torch_npu timm Pillow numpy safetensors

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch torch_npu timm Pillow numpy safetensors

# 从 ModelScope 下载 model.safetensors
wget "https://www.modelscope.cn/api/v1/models/timm/mobilenetv2_100.ra_in1k/repo?Revision=master&FilePath=model.safetensors" -O model.safetensors

python3 inference.py --device cpu

python3 inference.py --device npu

python3 compare_cpu_npu.py

import torch
import torch_npu
from timm import create_model
from safetensors.torch import load_file
from PIL import Image
from timm.data import create_transform, resolve_data_config

model = create_model('mobilenetv2_100', pretrained=False)
model.eval()
state_dict = load_file('model.safetensors')
model.load_state_dict(state_dict, strict=False)
model = model.to('npu:0')

img = Image.open('test.jpg').convert('RGB')
cfg = resolve_data_config({}, model=create_model('mobilenetv2_100', pretrained=False))
transform = create_transform(
    input_size=224, is_training=False,
    mean=cfg.get('mean'), std=cfg.get('std'),
    interpolation='bilinear')
input_tensor = transform(img).unsqueeze(0).to('npu:0')

with torch.no_grad():
    output = model(input_tensor)
probs = torch.nn.functional.softmax(output[0].cpu(), dim=0)
top5 = torch.topk(probs, k=5)
for i in range(5):
    print(f'Top {i+1}: class={top5.indices[i].item()}, prob={top5.values[i].item():.6f}')