Pytorch：当批量大小大于 1 时，如何绘制分割任务的预测输出？答案

【问题标题】：Pytorch: How plot the prediction output from segmentation task when batch size is larger than 1?Pytorch：当批量大小大于 1 时，如何绘制分割任务的预测输出？
【发布时间】：2019-05-03 09:52:31
【问题描述】：

我对分割主题中不同批次的绘图输出有疑问和疑问。

下面的 sn-p 绘制了每个类的概率和预测输出。

我确定概率图正在绘制一批，但不确定当我得到 torch.argmax(outputs, 1) 时的预测。我是否绘制了一批的 argmax，而网络的输出大小为 [10,4,256,256]。

另外，我想知道如何在批量大小为 10 时绘制所有批次的预测。

outputs = model(t_image)

fig, (ax1, ax2, ax3, ax4, ax5) = plt.subplots(nrows=1, ncols=5, sharex=True, sharey=True, figsize=(6,6))

img1 = ax1.imshow(torch.exp(outputs[0,0,:,:]).detach().cpu(), cmap = 'jet')
ax1.set_title("prob class 0")

img2 = ax2.imshow(torch.exp(outputs[0,1,:,:]).detach().cpu(), cmap = 'jet')
ax2.set_title("prob class 1")

img3 = ax3.imshow(torch.exp(outputs[0,2,:,:]).detach().cpu(), cmap = 'jet')
ax3.set_title("prob class 2")

img4 = ax4.imshow(torch.exp(outputs[0,3,:,:]).detach().cpu(), cmap = 'jet')
ax4.set_title("prob class 3")

img5 = ax5.imshow(torch.argmax(outputs, 1).detach().cpu().squeeze(), cmap = 'jet')
ax5.set_title("predicted")

【问题讨论】：

请说明您的尺寸相当于什么，即图像高度/宽度，以及您期望的输入/输出类型。

标签： python matplotlib pytorch prediction argmax

【解决方案1】：

不确定你在问什么。假设您使用的是 NCHW 数据布局，您的输出是每批 10 个样本、4 个通道（每个通道用于不同的类别）和 256x256 分辨率，那么前 4 个图表绘制了四个类别的类别分数。

对于第 5 个图，您的 torch.argmax(outputs, 1).detach().cpu().squeeze() 会给您一个 10x256x256 的图像，这是该批次中所有 10 个图像的类预测结果，而 matplotlib 无法正确地直接绘制它。所以你会想做torch.argmax(outputs[0,:,:,:], 0).detach().cpu().squeeze()，它会得到一个256x256的地图，你可以画出来。

由于结果的范围是从 0 到 3，代表 4 个类别（并且可能显示为非常暗淡的图像），通常人们会使用调色板为图着色。提供了一个示例here，与示例中的cityscapes_map[p] 行类似。

要绘制所有 10 个，为什么不写一个 for 循环：

for i in range(outputs.size(0)):
    # do whatever you do with outputs[i, ...]
    # ...
    plt.show()

并逐个检查批次中的每个结果。如果您的屏幕足够大，您还可以选择在子图中设置 10 行。

【讨论】：