语义分割入门系列之 U-Net

U-Net论文解读及代码解读

U-Net: Convolutional Networks for Biomedical Image Segmentation

UNet是基于FCN改良而来的网络，目的是用于医疗图像分割（也是医疗图像分割的基石）
网络结构如下图所示，整个网络包括两个部分，即左边的收缩路径和右边的扩张路径，收缩路径是经典的卷积神经网络模式，不断下采样并且加厚通道数，每次下采样的时候通道数翻倍。而与FCN不同的是，在右侧的扩张路径上，从最小分辨率的特征图开始，每次上采样的时候都进行2x2卷积，将通道数降为原来的一半，然后与左边的特征图进行concatenation（通道维度上的拼接），这里与FCN有所不同，FCN中特征之间的融合方式是summation。
细心的话可以发现，图中每次卷积特征图的尺寸都会减小2，当不对图片进行padding的时候，特征图尺寸就会不断缩小，最外面的一层边会被舍弃，所以把左边的特征图拼接到右边的时候需要裁剪出中间的特征图与右边的concat大小才对的上。

最后，由一个1x1的卷积将特征图上的通道维度转化为输出的分割标签，可见，每个图片的分割结果是比输入图片小一个恒定宽度的，如图，输入蓝色框的图片，输出的预测结果只有黄色框的部分，这是由于需要其余部分作为上下文语义信息帮助分割。

这里不禁有一个问题要问，这样最后的分割图片不就不完整了么？这就依靠文中提出的overlap-tile策略，对于UNet，输入任意大的图片，都可以把它想如图所示的这样分成很多小块，分别分割后再拼接起来实现无缝分割（这么做主要是由于一些医学图像分辨率较大），而确实的边缘部分则是用镜像，在原图上镜像扩展后，边缘部分就成了中心部分，这样把镜像出的一部分作为蓝框的输入，就能够得到边缘的结果。

导致这个麻烦的主要原因就是网络中的卷积都没有padding，不加padding的原因可能是因为加入padding后，会对边缘部分造成一定的误差，而随着语义信息的加深，padding积累的误差会越来越大。

Pytorch代码详解：
代码来自：https://github.com/milesial/Pytorch-UNet.git

import torch
import torch.nn as nn
import torch.nn.functional as F


class DoubleConv(nn.Module):  
"""两个连续卷积，这里padding是1，能够保持分辨率，按照论文的思路padding应该是0"""
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        return self.double_conv(x)

class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            self.up = nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)

        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        
        diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
        diffX = torch.tensor([x2.size()[3] - x1.size()[3]])
     """上采样操作的这一步需要注意，这里得到扩张通道的尺寸和左边收缩通道的特征图尺寸
        之间的差值，这个差值的一半作为padding的值添加到右侧特征图边缘上，使其尺寸与下采样过程中的
        特征图尺寸相同，然后再concat到一起做卷积，这里与论文有所不同，论文是没有padding的，
        左侧较大的特征图需要裁剪过后再与右侧较小的特征图concat """
        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])

        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

class UNet(nn.Module):
"""UNet 结构部分，四个下采样和四个上采样，上采样时，如self.up2(x, x3)，结合了下采样过程中的特征图"""
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 512)
        self.up1 = Up(1024, 256, bilinear)
        self.up2 = Up(512, 128, bilinear)
        self.up3 = Up(256, 64, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits