mobilenetv4使用笔记

demo1.py

import time

from urllib.request import urlopen
import torch
from PIL import Image
import timm

# 打开图像
img = Image.open(urlopen('https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'))

# 创建模型，设置 pretrained=False 以避免从网络加载预训练权重
model = timm.create_model('mobilenetv4_conv_large.e600_r384_in1k', pretrained=False)

# 手动加载本地的预训练权重
pretrained_weights_path = 'models/pytorch_model.bin'
model.load_state_dict(torch.load(pretrained_weights_path))

# 设置模型为评估模式
model = model.eval()
model = model.cuda()

# 获取模型特定的变换（归一化、调整大小等）
data_config = timm.data.resolve_model_data_config(model)
transforms = timm.data.create_transform(**data_config, is_training=False)

for i in range(10):
    img = transforms(img).cuda()
    start=time.time()
    output = model(img.unsqueeze(0))  # unsqueeze 将单张图片扩展为批量大小为1
    print(img.shape,output.shape,'time', time.time()-start)

# 获取 top 5 结果
top5_probabilities, top5_class_indices = torch.topk(output.softmax(dim=1) * 100, k=5)

# 打印结果
print(top5_probabilities)
print(top5_class_indices)

gpu 384 平均7ms