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

Question

对segmentation subject中不同batch的plot输出有疑惑和疑问

下面的代码片段绘制了每个 class 的概率和预测输出。

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

此外，我想知道当我的批次大小为 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")

Answer 1

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

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

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

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

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

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

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

Pytorch: How plot the prediction output from segmentation task when batch size is larger than 1?

python

matplotlib

pytorch

prediction

argmax