19.predict

预测阶段大致步骤：

打开图片：[指定路径]，[PIL打开]

转换图片：[RGB],[resize,totensor]

调用训练模型：[torch.load]

预测

输出结果:[argmax(1)]

要注意的是：

CUDA训练好的模型，预测时：1，device为cpu，调用时映射到cpu上

                        2，device为gpu，将输入图片转化为CUDA类型

# 1
model = torch.load("pretrained_demo_10.pth",map_location=torch.device("cpu"))
# 2
device =  torch.device("cuda" if torch.cuda.is_available() else "cpu")
# img = transform(img)后
img = img.to(device)
# or 
output = model(img.to(device))

预测结果为：dog √

code

import torch
import torchvision
from PIL import Image
# GPU训练的模型，测试时也要使用GPU
device =  torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 打开测试图片
image_path = "./img/dog.png"
img = Image.open(image_path)
# 很多截图软件截图是三通道外加一个透明通道， 所以转换一下
img = img.convert("RGB")
#  resize and totensor
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
                                            torchvision.transforms.ToTensor()])
img = transform(img)
# img = img.to(device)必须加在这里 或使用output = model(img.to(device))
# >>torch.Size([3, 32, 32])
# 调用训练好的模型
model = torch.load("pretrained_demo_10.pth")
#model = torch.load("pretrained_demo_10.pth",map_location=torch.device("cpu"))
# print(model)
# 模型接收是4维 所以reshape增加一个batch_size维度
img = torch.reshape(img,(1,3,32,32))
# 测试
model.eval()
with torch.no_grad():
    output = model(img.to(device))
    #output = model(img)
# 每个类别的得分
print(output)
# 输出预测类别
print(output.argmax(1))