一、选取的图片
全部代码依托于该博客
二、代码(调用训练好的模型)
python">import torch
import torchvision
from PIL import Image
from model import *img_path = "dog.png"
image = Image.open(img_path)print(image.size)transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),torchvision.transforms.ToTensor()])
image = transform(image)
image = torch.reshape(image,(1,3,32,32))print(image.shape)#更换pth文件
model = torch.load('model.pth', map_location=torch.device('cpu'))
print(model)model.eval()
with torch.no_grad():output = model(image)
print(output.shape)
print(output)
print(torch.argmax(output,dim=1))
输出:
python">(287, 251)
torch.Size([1, 3, 32, 32])
Mary((model1): Sequential((0): Conv2d(3, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))(1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(2): Conv2d(32, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))(3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(4): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))(5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(6): Flatten(start_dim=1, end_dim=-1)(7): Linear(in_features=1024, out_features=64, bias=True)(8): Linear(in_features=64, out_features=10, bias=True))
)
torch.Size([1, 10])
tensor([[-2.9548, -3.5498, 0.3010, 3.4690, -0.8843, 4.8731, 2.5652, -0.5132,-1.0659, -2.9456]])
tensor([5])
模型预测的图片类别为 5 ,可以看到模型预测的结果是正确的