【pytorch图像分类】批量图像窗口显示、测试结果可视化

深度学习模型搭建好之后，想在测试时显示每张图像的预测类别和真实类别。查资料后，我找到了这个方法，每张图片都能显示出来，效果如下：
（这里我的4种类别分别用的0、1、2、3代替）

在程序中加入这个函数

def imshow(inp, title, ylabel):
    """Imshow for Tensor."""
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)
    plt.show()
    plt.ylabel('GroundTruth: {}'.format(ylabel))
    plt.title('predicted: {}'.format(title))

再在测试集循环的过程中加入几行代码即可

 if j%4==0:
             plt.figure()
             j=j%4
        plt.subplot(2, 2, j + 1)
        imshow(out, title=[classes[predicted]], ylabel=[classes[labels]])

整体的代码如下：

from __future__ import print_function, division
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import torch.nn as nn
import torch
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import time
import os
import copy
from torchvision import models
import matplotlib.pyplot as plt
import torchvision


test_transform=transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

test_dataset = torchvision.datasets.ImageFolder(root='maize/test',transform=test_transform)
test_loader = DataLoader(test_dataset,batch_size=1, shuffle=True,num_workers=0)#num_workers：使用多进程加载的进程数，0代表不使用多进程

classes = test_dataset.classes

model = torch.load('model.pkl')

def imshow(inp, title, ylabel):
    """Imshow for Tensor."""
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)
    plt.show()
    plt.ylabel('GroundTruth: {}'.format(ylabel))
    plt.title('predicted: {}'.format(title))


correct = 0
total = 0
i=1
j=0
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        out = torchvision.utils.make_grid(images)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        print(i,'.Predicted:', ''.join('%5s' % classes[predicted] ),'  GroundTruth:',''.join('%5s' % classes[labels] ))
        if j%4==0:#设置每个窗口显示4张
             plt.figure()
             j=j%4
        plt.subplot(2, 2, j + 1)
        imshow(out, title=[classes[predicted]], ylabel=[classes[labels]])
        j=j+1
        i=i+1
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print('Accuracy of the network on the test images: %d %%' % (100 * correct / total))

测试数据集为文件夹读取，具体可看这篇文章
文件夹读取数据集
希望对你有帮助！