[Pytorch]语义分割任务分类的实现

[Pytorch]语义分割任务分类的实现

假如我们定义了一个网络用于语义分割任务，这个网络简称为model()
语义分割任务要做的是：

• 对于一个图片输入input，大小为（B，C，W， H)，其中B为batchsize，C为通道数channel，我们的网络输出一个W，H保持为原图大小的输出output，而通道数改为分类类别数num_class，目的是对于原图每一个像素位置都进行分类，给每一个像素都分出一个预测类别标签。
• 所以语义分割网络的输出output的尺寸应该为（B，N，W，H）,其中N为num_class

对于输出outputs，（W，H）范围内的每一个位置在通道N的第i个通道位置的值，都代表了将该坐标（x，y）的像素分类为第i类的类别得分。
例如如下初始化了一个tensor来模拟对应的output：

out = torch.tensor([[[1, 2],[4, 5],],[[1, 7],[0, 3]],[[2, 1],[6, 2]]
])

其中中间:
[1, 7]
[0, 3]
这个矩阵就代表，在这个2x2的图片中，将四个位置的像素预测成第1类（类别从0开始）的类别得分分别为：1，7，0，3
最终预测出每一个位置的类别应该为所有通道N中取最大的那一个通道所对应的类别：

# 模拟batchsize为1的情况
out = out.unsqueeze(0)
print(out.shape)

输出：torch.Size([1, 3, 2, 2])

使用tensor.max(dim)方法统计对应维度的最大值，我们这里所要统计的是通道维度上的最大值，所以使用以下方法：

out = out.max(1)

输出：torch.return_types.max(
values=tensor([[[2, 7],
[6, 5]]]),
indices=tensor([[[2, 1],
[2, 0]]]))

out.max(1)方法返回了两个列表，其中第一个列表是每一个位置对应通道中的最大值，而第二个列表返回的是对应通道的索引，也就是对应的类别，我们在实际预测中是以分类为目的所以取第二个列表，改为：

out = out.max(1)[1]

输出：tensor([[[2, 1],
[2, 0]]])

这样一来，矩阵
[2, 1]
[2, 0]
就是最终每一个像素的分类结果