Pytorch学习--神经网络--完整的模型验证套路

一、选取的图片

全部代码依托于该博客

二、代码（调用训练好的模型）

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))

输出：

(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 ，可以看到模型预测的结果是正确的