预测模型

import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet

# 图片处理大小核标准化
transform = transforms.Compose(
    [transforms.Resize((32, 32)),
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# 实例化 LeNet
net = LeNet()
# 载入权重文件
net.load_state_dict(torch.load('Lenet.pth'))

im = Image.open('1.jpg')      # 保存到目录下一张图片用来预测
im = transforms(im)  # 预处理为 [C, H, W]
im = torch.unsqueeze(im, dim=0)  # 增加新维度 tensor格式，用来网络中传播[N, C, H, W]

with torch.no_grad():
    outputs = net(im)
    predict = torch.max(outputs, dim=1)[1].data.numpy()
    # 或者用softmax函数
    # predict = torch.softmax(outputs, dim=1) 
print(classes[int(predict)])