【项目实战】pytorch载入训练好的模型并进行可视化模型预测绘图

main函数载入模型，加载图片，输出结果：

if __name__ == '__main__':
  image =  Image.open(r"C:\Users\pic\test\he_5.jpg")
    image =transform(image).unsqueeze(0)
    modelme = torch.load('modefresnet.pkl')
    modelme.eval() #表示将模型转变为evaluation（测试）模式，这样就可以排除BN和Dropout对测试的干扰。
     visualize_model(modelme)
    outputs = modelme(image)
    _, predict = torch.max(outputs.data, 1)
        for j in range(image.size()[0]):
     print('predicted: {}'.format(class_names[predict[j]]))

对图片的统一处理transform:

transform=transforms.Compose([
            transforms.Resize(224),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485,0.456,0.406],
                                 std=[0.229,0.224,0.225])
                            ])

对于预测结果进行可视化的函数：

def visualize_model(model, num_images=6):
 was_training = model.training
 model.eval()
 images_so_far = 0
 fig = plt.figure()
 with torch.no_grad():

     #for i, (inputs, labels) in enumerate(dataloaders['val']):
     for i, (inputs, labels) in enumerate(testloder):
       outputs = model(inputs)
       _, preds = torch.max(outputs, 1)
       
       for j in range(inputs.size()[0]):

           images_so_far += 1

           ax = plt.subplot(num_images // 2, 2, images_so_far)

           ax.axis('off')

           ax.set_title('predicted: {}'.format(class_names[preds[j]]))

           imshow(inputs.cpu().data[j])


           if images_so_far == num_images:

            model.train(mode=was_training)
            plt.show()

            return

     model.train(mode=was_training)

载入一新的图片数据集：

data_dir =os.getcwd() + '\\data\\'
dataloadertest =datasets.ImageFolder(os.path.join(data_dir, "tt"),transform=transforms.Compose([
            transforms.Resize(224),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485,0.456,0.406],
                                 std=[0.229,0.224,0.225])
                            ]) )
testloder = torch.utils.data.DataLoader(dataloadertest,batch_size = 4,shuffle = True)

目录结构:

其中要注意传入的图片的预处理：
image = Image.open(r"C:\Users\pic\test\he_5.jpg")
image =transform(image).unsqueeze(0)
需为PIL格式，且需先进行转化才能传入模型。

结果:

经测试之后不论是传入单张图片还是一个新数据集结果均符合预期。