第一步：准备测试图片和测试模型

准备好一张需要测试的图片和训练好的测试模型，放到测试文件夹中

第二步：测试代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from PIL import Image
from torchvision import transforms
import torch
from torch.autograd import Variable
#加载测试模型
model = torch.load('C:/Users/Arthur/Desktop/test/model.pth')
#加载测试图片
test_img = Image.open('C:/Users/Arthur/Desktop/test/12.png')
#图片变换
transform = transforms.Compose([transforms.Resize((224,224))
                                       ,transforms.ToTensor()
                                       ,transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
test_img = transform(test_img)
#增加一维度
inputs =Variable(torch.unsqueeze(test_img, dim=0).float(), requires_grad=False).cuda()
outputs = model(inputs)
#经过sofamax层后输出概率
probability = torch.nn.functional.softmax(outputs,dim=1)
#找出最大概率和相应的标签
max_value,index = torch.max(probability,1)
print(probability,max_value,index)

注：

  1. 测试图片的尺寸需要和训练尺寸相同，并且也需要进行张量变换（ToTensor）和归一化（Normalize）。
  2. 由于pytorch要求的输入的维度为[batch_size, channels, width, height]，而一个样本的维度为[channels, width, height]，此时因为我们需要测试一张图片，所以用unsqueeze()增加一个维度变为[1, channels, width, height]。
  3. index只能给你一个数字，它其实代表的是第几个类别，如果想知道具体类别名称，可以打开训练集的文件夹，那么0就是第一个类别，1就是第二个类别…….以此类推。比如以下是我的训练集。
    训练集文件夹
    那么mn就是第0个类别，mp就是第一个类别。