mobilenetv4_conv_blur_medium.e500_r224_in1k 昇腾 NPU 适配

pip install torch torch_npu timm Pillow numpy safetensors modelscope

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

python -c "from modelscope.hub.snapshot_download import snapshot_download; snapshot_download('timm/mobilenetv4_conv_blur_medium.e500_r224_in1k')"

python inference.py

python 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('mobilenetv4_conv_blur_medium', pretrained=False, num_classes=1000)
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')
config = resolve_data_config({{}}, model=create_model('mobilenetv4_conv_blur_medium', pretrained=False, num_classes=1000))
transform = create_transform(
    input_size=224, is_training=False,
    mean=config.get('mean', (0.485, 0.456, 0.406)),
    std=config.get('std', (0.229, 0.224, 0.225)),
    interpolation=config.get('interpolation', 'bicubic'))
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)
print(f'Top-1 class: {probs.argmax().item()}')

# NPU 推理
python3 inference.py --device npu

# CPU 推理
python3 inference.py --device cpu