想想这样一个场景,我们训练了一个神经网络,输入一些信息,这个网络可以根据信息为我们生成相关图片。
这些图片并不是一张,而是多张,我们想把这些图片一次全部显示出来,而不是一张一张的显示(这种似乎太浪费时间了).
先看示意图:
借助torchvision库,我们可以很容易实现此功能。简单来看,代码是这样的:
先看一下IMG类型
>>> IMG.size()
torch.Size([16, 3, 640, 640])
>>> type(IMG)
<class 'torch.Tensor'>
grid=torchvision.utils.make_grid(IMG,nrow=4)#nrow指定列数
plt.imshow(np.transpose(grid,(1,2,0)))#将张量转换为矩阵
plt.show()#绘图
torchvision.utils.save_image(grid,'out.jpg')#保存网格图片
上述代码已经足够完成任务了,当然,torchvision.utils.make_grid方法可接受许多参数,想了解更多,可见PyTorch官方文档。