文章目录
- [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]
就是最终每一个像素的分类结果