虽然 images,labels = next(iter(test_data)) 可以每次1batch获取,但是超级慢,不推荐。
import torch import platform if platform.system() == 'Windows': import matplotlib matplotlib.use('TkAgg') import matplotlib.pyplot as plt elif platform.system() == 'Linux': import matplotlib.pyplot as plt #.pt的加载方式 model = torch.load(r'best_001_1.000000_model.pt') device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) #模型转入GPU model.eval() #测试模式 # 查看前10batch的预测效果图 plt.figure(figsize=(12, 4)) # 创建窗口并指定大小(英寸) for index,(images,labels) in enumerate(test_data): images,labels = images.cuda(),labels.cuda() #样本转入GPU outputs = model(images) #预测1个batch _, predicted_labels = torch.max(outputs.data, 1) #显示1个batch的图像,title:预测标签-真实标签,预测错误title为红色 for i in range(batch): images_cpu = images.cpu() #GPU转到CPU来进行显示 img = images_cpu[i].squeeze(0).permute(1, 2, 0) #NCHW变为CHW再变为HWC img = (img+1)/2 #[-1,1]变为[0,1]。[0,1]或[0,255]都可以显示 plt.subplot(1,4,i+1) #1行4列来排布 plt.title('{}-{}'.format(classes[predicted_labels[i].item()], classes[labels[i].item()]), color='red' if classes[predicted_labels[i].item()]!=classes[labels[i].item()] else 'black')#1个batch里的图对应的标签为标题 plt.imshow(img) #显示图像 plt.axis('off') #不显示坐标轴 # plt.show() #显示窗体并阻塞 plt.ion() #显示窗口,不阻塞 plt.pause(0.1) #阻塞0.1s,其实是让窗口一直驻留显示 input('Press Enter to Continue {}'.format(index+1)) if index==9: plt.close() break
标签:plt,torch,展示,labels,matplotlib,Enter,images,import From: https://www.cnblogs.com/xixixing/p/18520490