本次测试输入 dog.png
# Coding by ajupyterfrom PIL import Imagefrom torch import nnimport torchimport torchvisionclass 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 outputmodel = 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