如何基于PyTorch计算语义分割中的Top3 IoU？

给定神经网络输出outputs（形状为(N, C, H, W)）和真实标签label（形状为(N, H, W)），其中：

• N：批量大小
• C：类别总数
• H、W：特征图的高和宽
  label中每个元素的取值范围为[0, C-1]，代表对应像素的真实类别。

要计算Top3 IoU，我们需要生成仅标记真实类别是否落在模型Top3预测结果中的One-Hot掩码。以下是正确的实现方式：

import torch

N, C, H, W = 1, 4, 2, 2
outputs = torch.rand((N, C, H, W))
label = torch.arange(C).reshape(N, H, W)
_, index = torch.topk(outputs, k=3, dim=1)
top3 = torch.zeros((N, C, H, W))

for i in range(N):
    for j in range(H):
        for k in range(W):
            c = label[i, j, k]
            if c in index[i, :, j, k]:
                top3[i, c, j, k] = 1

示例输出： #

• outputs：

tensor([[[[0.8002, 0.6733],
          [0.7034, 0.5039]],

         [[0.8401, 0.9226],
          [0.7963, 0.6157]],

         [[0.1063, 0.0310],
          [0.2489, 0.9920]],

         [[0.8279, 0.9109],
          [0.4737, 0.2299]]]])

• label：

tensor([[[0, 1],
         [2, 3]]])

• index（Top3预测的类别索引）：

tensor([[[[1, 1],
          [1, 2]],

         [[3, 3],
          [0, 1]],

         [[0, 0],
          [3, 0]]]])

• top3（正确的One-Hot掩码）：

tensor([[[[1., 0.],
          [0., 0.]],

         [[0., 1.],
          [0., 0.]],

         [[0., 0.],
          [0., 0.]],

         [[0., 0.],
          [0., 0.]]]])

生成top3后，即可用它来计算Top3 IoU。

错误实现示例 #

有一类基于one_hot的实现常用于计算Top3像素准确率，但生成的掩码会包含大量错误的1值，完全无法用于IoU计算：

expand = torch.nn.functional.one_hot(index)
top3 = expand.transpose(1, 4).sum(dim=4) 

其输出结果为：

tensor([[[[0, 1],
          [0, 0]],

         [[1, 0],
          [1, 1]],

         [[1, 1],
          [1, 1]],

         [[1, 1],
          [1, 1]]]])

可以看到该结果与正确的top3掩码差异极大，错误地将多个类别标记为1，无法满足IoU计算的需求。

内容的提问来源于stack exchange，提问作者fatpanda2049