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

小灰灰 2022-11-05 10:52 253阅读 0赞

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)

目录结构:  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQxMzU4NTc0_size_16_color_FFFFFF_t_70]  
其中要注意传入的图片的预处理：  
image = Image.open(r"C:\\Users\\pic\\test\\he\_5.jpg")  
image =transform(image).unsqueeze(0)  
需为PIL格式，且需先进行转化才能传入模型。

结果:  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQxMzU4NTc0_size_16_color_FFFFFF_t_70 1]

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQxMzU4NTc0_size_16_color_FFFFFF_t_70 2]  
经测试之后不论是传入单张图片还是一个新数据集结果均符合预期。

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQxMzU4NTc0_size_16_color_FFFFFF_t_70]: /images/20221023/cc059c51c30d4b0bada21e3b73358b76.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQxMzU4NTc0_size_16_color_FFFFFF_t_70 1]: /images/20221023/5b25ef9e5ae240f98b87b2e4d349f3ea.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQxMzU4NTc0_size_16_color_FFFFFF_t_70 2]: /images/20221023/b5e0758b452647898bfa2fb8c3a0876f.png