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)])
预测模型
最新推荐文章于 2024-01-03 17:16:21 发布