b站小土堆pytorch学习记录——P32 完整的模型验证套路

CIFAR10数据集中不同类别对应的target：

代码：

import torchimport torchvisionfrom PIL import Imagefrom torch import nnimg_path = "plane.png"img = Image.open(img_path)transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),                                           torchvision.transforms.ToTensor()])img = transform(img)class Guodong(nn.Module):    def __init__(self):        super(Guodong,self).__init__()        self.module = nn.Sequential(            nn.Conv2d(3, 32, 5, 1, 2),            nn.MaxPool2d(2),            nn.Conv2d(32, 32, 5, 1, 2),            nn.MaxPool2d(2),            nn.Conv2d(32, 64, 5, 1, 2),            nn.MaxPool2d(2),            nn.Flatten(),            nn.Linear(64*4*4, 64),            nn.Linear(64, 10)        )    def forward(self, x):        x = self.module(x)        return xmodel = torch.load("guodong_9.pth")//这是我已经训练好了的模型print(model)img = torch.reshape(img, (1, 3, 32, 32))model.eval()with torch.no_grad():    output = model(img)print(output)print(output.argmax(1))

plane.png图片如下：

训练模型的代码见上一篇博客：

模型训练套路

运行结果：

可以看到output最大arg为0，而0正好对应airplane，表示该分类成功