使用torchvision中的模型进行预测

from PIL import Image
import matplotlib.pyplot as plt
import json
import numpy as np

import torch
import torch.nn.functional as F
from torchvision import models, transforms

# 构造模型，读取权重
# model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)  # 使用预训练模型,模型自动下载到C:\Users\GaryLau\.cache\torch\hub\checkpoints
model = models.resnet18()               # 不使用预训练模型
model.load_state_dict(torch.load('resnet18-f37072fd.pth'))   # 加载权重
model.eval()   # 设置为eval模式

transform = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(),
                                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

# 读取图片，预处理
im = Image.open('dog.1.jpg')
transformed_img = transform(im)
input_img = transformed_img.unsqueeze(0)  # 增加批次维度

# 模型推理
output = model(input_img)
output_softmax = F.softmax(output, dim=1)
print(output_softmax.size())     # torch.Size([1, 1000])

# 从预测结果中取出top5
prediction_score, pred_label_idx = torch.topk(output_softmax, 5)
prediction_score_ = prediction_score.squeeze().detach().numpy()
pred_label_idx_ = pred_label_idx.squeeze().numpy()   # 获取预测结果的标签索引
classes = json.load(open("imagenet1000.json"))       # 读取标签文件
class_name = classes[pred_label_idx_[0]]             # 根据预测的标签索引得到预测标签名称
print("Pred Class index: {}, Pred Class name: {}".format(pred_label_idx_[0] , class_name))

# 可视化处理，创建一个一行两列的子图
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 9))
fig.sca(ax1)       # 设置第一个轴是ax1
ax1.imshow(im)     # 第一个子图显示原始要预测的图
# 设置第二个图是预测的解雇，按概率取出top5
barlst = ax2.bar(range(5), [i for i in prediction_score_])
barlst[0].set_color('red')   # 设置top1的颜色是红色
# 预测结果前5的柱状图
plt.sca(ax2)
plt.ylim([0, 1.1])
# 垂直显示前5的标签
plt.xticks(range(5), [classes[i] for i in pred_label_idx_], rotation='vertical')
fig.subplots_adjust(bottom=0.2)
plt.show()

print('Done.')

输出：

torch.Size([1, 1000])
Pred Class index: 209, Pred Class name: Chesapeake Bay retriever
Done.

结果图：

原图：
dog.1.jpg：

上面使用的是resnet18模型，可以根据模型的规模、精度要求挑选更为合适的模型。