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