Pytorch预测

本次测试输入 dog.png

# Coding by ajupyter
from PIL import Image
from torch import nn
import torch
import torchvision


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 64),
            nn.Linear(64, 10),
        )

    def forward(self, x):
        output = self.model(x)
        return output


model = Model()
model.load_state_dict(torch.load('cifa10_model-epoch19-test_loss699.5733557939529'))

model.eval()
with torch.no_grad():

    image = Image.open('dog.png')
    image = image.convert('RGB')  # 适配jpg和png jpg是四通道：rgb+透明度通道

    tool = torchvision.transforms.Compose([
        torchvision.transforms.Resize(size=(32, 32)),
        torchvision.transforms.ToTensor()
    ])
    input = tool(image)
    input = input.reshape((1, 3, 32, 32))
    print(input.shape)  # torch.Size([3, 32, 32])
    res = model(input)
    print(res)
    print(res.argmax(1))  # 横向比较
    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    print(f'res:{classes[res.argmax(1).item()]}')  # only one element tensors can be converted to Python scalars