Pytorch实现FCN中对损失函数的理解

代码：https://github.com/bat67/pytorch-FCN-easiest-demo

FCN网络的输出

输入网络的图片大小为H×W，通道数为3（RGB图像）。
而经过卷积和上采样过程，输出图片大小不变，仍为H×W，通道数为分割的类别数C。
在Pytorch中输出的shape为torch.Size([C, H, W])

FCN网络的损失函数

criterion = nn.BCELoss().to(device)
loss = criterion(output, mask)

可以看出，损失函数是计算output（即FCN网络输出）与mask之间的交叉熵损失。
我们需要理解的是，mask如何转换成与output相同的shape，即torch.Size([C, H, W])，这样两者之间才能计算损失。

对mask数据的处理

以分割bag图片为例，分割类别为2类（C=2），分别为bag和background，如图所示，左侧为原始图片，右侧为图片对应的mask。

mask = cv2.imread('path', 0)
mask = cv2.resize(mask, (10, 10))
mask = mask / 255
mask = mask.astype('uint8')
mask = onehot(mask, 2)
mask = mask.transpose(2, 0, 1)
mask = torch.FloatTensor(mask)

• 第一步，读取mask图片，注意这里imread函数的flag=0，表示为灰度图，此时图片大小为H×W，通道数为1。
  这里我们为了便于显示mask中的值，将它大小改为10×10（原始代码中是160×160）
  

mask = cv2.imread('path', 0)
mask = cv2.resize(mask, (10, 10))

• 第二步，将mask归一化，这时像素值变为0-1之间的浮点数
  

mask = mask / 255

• 第三步，将浮点数转换为整型，这时像素值变为0或1，此时我们就能理解了，像素值为0的点代表background，而像素值为1的点代表bag。
  

mask = mask.astype('uint8')

• 第四步，将mask转换成one-hot编码格式
  

mask = onehot(mask, 2)

onehot函数：

def onehot(data, n):
    buf = np.zeros(data.shape + (n, ))
    nmsk = np.arange(data.size) * n + data.ravel()
    buf.ravel()[nmsk - 1] = 1
    return buf

这时，mask的大小从10×10变成了10×10×2，原本的像素值由0 / 1，变成了[0, 1] / [1, 0]，也就是将通道数从1变成了2。

• 最后一步，将numpy.ndarray格式的mask转换成Tensor。
  

mask = mask.transpose(2, 0, 1)  # 10×10×2 → 2×10×10
mask = torch.FloatTensor(mask)  # numpy.ndarray → torch.Tensor

此时，mask的shape就转换为了torch.Size([C, H, W])，即torch.Size([2, 10, 10])，与FCN网络的output大小相同，二者之间就可以计算损失函数了。