Unet网络解析

1 Unet网络概述

论文名称：U-Net: Convolutional Networks for Biomedical Image Segmentation

发表会议及时间 ：MICCA ( 国际医学图像计算和 计算机辅 助干预会 议 ) 2 0 1 5

Unet提出的初衷是为了解决医学图像分割的问题。

Unet网络非常的简单，前半部分就是特征提取，后半部分是上采样。在一些文献中把这种结构叫做编码器-解码器结构，由于网络的整体结构是一个大些的英文字母U，所以叫做U-net。其实可以将图像->高语义feature map的过程看成编码器，高语义->像素级别的分类score map的过程看作解码器

• Encoder：左半部分，由两个3x3的卷积层（RELU）再加上一个2x2的maxpooling层组成一个下采样的模块；
• Decoder：右半部分，由一个上采样的卷积层（反卷积层），特征拼接concat，两个3x3的卷积层，非线性ReLU层；

在当时，Unet相比更早提出的FCN网络，使用拼接来作为特征图的融合方式。

• FCN是通过特征图对应像素值的相加来融合特征的；
• U-net通过通道数的拼接，这样可以形成更厚的特征，当然这样会更佳消耗显存；

2 Unet与FCN网络的区别

U-Net和FCN非常的相似，U-Net比FCN稍晚提出来，但都发表在2015年，和FCN相比，U-Net的第一个特点是完全对称，也就是左边和右边是很类似的，而FCN的decoder相对简单。第二个区别就是skip connection，FCN用的是加操作（summation），U-Net用的是叠操作（concatenation）。这些都是细节，重点是它们的结构用了一个比较经典的思路，也就是编码和解码（encoder-decoder)结构。其实可以将图像->高语义feature map的过程看成编码器，高语义->像素级别的分类score map的过程看作解码器

此外, 由于UNet也和FCN一样, 是全卷积形式, 没有全连接层(即没有固定图的尺寸),所以容易适应很多输入尺寸大小,但并不是所有的尺寸都可以,需要根据网络结构决定

3 为什么Unet在医疗图像分割种表现好

• 医疗影像语义较为简单、结构固定。因此语义信息相比自动驾驶等较为单一，因此并不需要去筛选过滤无用的信息。医疗影像的所有特征都很重要，因此低级特征和高级语义特征都很重要，所以U型结构的skip connection结构（特征拼接）更好派上用场
• 医学影像的数据较少，获取难度大，数据量可能只有几百甚至不到100，因此如果使用大型的网络例如DeepLabv3+等模型，很容易过拟合。大型网络的优点是更强的图像表述能力，而较为简单、数量少的医学影像并没有那么多的内容需要表述，因此也有人发现在小数量级中，分割的SOTA模型与轻量的Unet并没有优势
• 医学影像往往是多模态的。比方说ISLES脑梗竞赛中，官方提供了CBF，MTT，CBV等多中模态的数据（这一点听不懂也无妨）。因此医学影像任务中，往往需要自己设计网络去提取不同的模态特征，因此轻量结构简单的Unet可以有更大的操作空间。

因此，大多数医疗影像语义分割任务都会首先用Unet作为baseline

4 Unet网络结构

Unet网络是建立在FCN网络基础上的，它的网络架构如下图所示，总体来说与FCN思路非常类似。这里需要注意的是，U-Net的输入大小是572x572，但是输出却是388x388，按理说它们应该相等（因为图像分割相当于逐像素进行分类，所以要求输入图像和输出图像大小一致），但是为什么这里的输入尺寸要比输出尺寸大呢？那是因为下图这个结构图是当年论文作者绘制的，该作者对输入图像的边缘进行了镜像填充，通过镜像填充将边界区域进行扩大，这样可以给模型提供更多信息来完成模型的分割。

按照论文中的解释，镜像填充的原因是：因为图像 的边界的外面是空白的，没有其它有效像素，而我们预测图像中的像素类别时往往需要参考它的周围像素作为上下文信息，这样才能保持分割的准确性，为了预测边界像素，论文对边界区域进行镜像，来补全边界周围缺失的内容，然后进行预测。这种策略叫做"overlap-tile"

这里的输入是单通道的原因是因为输入图片是灰度图，而输出是两通道是因为这里是对像素进行二分类（前景和背景），所以输出通道是2

整个网络由编码部分（左） 和 解码部分（右）组成，类似于一个大大的U字母，具体介绍如下：

1、编码部分是典型的卷积网络架构：

• 它主要的作用是进行特征提取
• 架构中含有着一种重复结构，每次重复中都有2个 3 x 3卷积层、非线性ReLU层和一个 2 x 2 max pooling层（stride为2）。（图中的蓝箭头、红箭头，没画ReLu）
• 每一次下采样后我们都把特征通道的数量加倍

2、解码部分也使用了类似的模式：

• 它主要的作用是进行上采样 (上采样可以让包含高级抽象特征的低分辨率图片在保留高级抽象特征的同时变为高分辨率)
• 架构中包含有一种重复结构，每次重复都有一个上采样的卷积层（反卷积层），特征拼接concat，两个3x3的卷积层，非线性ReLU层
• 每一步都首先使用反卷积(up-convolution)，每次使用反卷积都将特征通道数量减半，特征图大小加倍。（图中绿箭头）
• 反卷积过后，将反卷积的结果与编码部分中对应步骤的特征图拼接起来(concat)（也就是将深层特征与浅层特征进行融合，使得信息变得更丰富）。（白/蓝块）
• 编码部分中的特征图尺寸稍大，将其修剪过后进行拼接（这里是将两个特征图的尺寸调整一致后按通道数进行拼接）。（左边深蓝虚线部分就是要裁剪的部分，它对应右边的白色长方块部分）
• 对拼接后的map再进行2次3 x 3的卷积。（右侧蓝箭头）
• 最后一层的卷积核大小为1 x 1，将64通道的特征图转化为特定类别数量（分类数量）的结果。（图中青色箭头）

5 代码复现

下面使用pytorch框架对论文中的unet进行复现：

import torch.nn as nn
import torch

# 编码器(论文中称之为收缩路径)的基本单元
def contracting_block(in_channels, out_channels):
    block = torch.nn.Sequential(
        # 这里的卷积操作没有使用padding,所以每次卷积后图像的尺寸都会减少2个像素大小
        nn.Conv2d(kernel_size=(3, 3), in_channels=in_channels, out_channels=out_channels),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(),
        nn.Conv2d(kernel_size=(3, 3), in_channels=out_channels, out_channels=out_channels),
        nn.BatchNorm2d(out_channels),
        nn.ReLU()
    )
    return block

# 解码器（论文中称之为扩张路径）的基本单元
class expansive_block(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels):
        super(expansive_block, self).__init__()

        # 每进行一次反卷积，通道数减半，尺寸扩大2倍
        self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=(3, 3), stride=2, padding=1,
                                     output_padding=1)
        self.block = nn.Sequential(
            # 这里的卷积操作没有使用padding,所以每次卷积后图像的尺寸都会减少2个像素大小
            nn.Conv2d(kernel_size=(3, 3), in_channels=in_channels, out_channels=mid_channels),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(),
            nn.Conv2d(kernel_size=(3, 3), in_channels=mid_channels, out_channels=out_channels),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )

    def forward(self, e, d):
        d = self.up(d)
        # concat
        # e是来自编码器部分的特征图，d是来自解码器部分的特征图，它们的形状都是[B,C,H,W]
        diffY = e.size()[2] - d.size()[2]
        diffX = e.size()[3] - d.size()[3]
        # 裁剪时，先计算e与d在高和宽方向的差距diffY和diffX，然后对e高方向进行裁剪，具体方法是两边分别裁剪diffY的一半，
        # 最后对e宽方向进行裁剪，具体方法是两边分别裁剪diffX的一半，
        # 具体的裁剪过程见下图一
        e = e[:, :, diffY // 2:e.size()[2] - diffY // 2, diffX // 2:e.size()[3] - diffX // 2]
        cat = torch.cat([e, d], dim=1)  # 在特征通道上进行拼接
        out = self.block(cat)
        return out

# 最后的输出卷积层
def final_block(in_channels, out_channels):
    block = nn.Conv2d(kernel_size=(1, 1), in_channels=in_channels, out_channels=out_channels)
    return block

class UNet(nn.Module):

    def __init__(self, in_channel, out_channel):
        super(UNet, self).__init__()

        # 编码器 (Encode)
        self.conv_encode1 = contracting_block(in_channels=in_channel, out_channels=64)
        self.conv_pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv_encode2 = contracting_block(in_channels=64, out_channels=128)
        self.conv_pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv_encode3 = contracting_block(in_channels=128, out_channels=256)
        self.conv_pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv_encode4 = contracting_block(in_channels=256, out_channels=512)
        self.conv_pool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        # 编码器与解码器之间的过渡部分(Bottleneck)
        self.bottleneck = nn.Sequential(
            nn.Conv2d(kernel_size=(3, 3), in_channels=512, out_channels=1024),
            nn.BatchNorm2d(1024),
            nn.ReLU(),
            nn.Conv2d(kernel_size=(3, 3), in_channels=1024, out_channels=1024),
            nn.BatchNorm2d(1024),
            nn.ReLU()
        )

        # 解码器(Decode)
        self.conv_decode4 = expansive_block(1024, 512, 512)
        self.conv_decode3 = expansive_block(512, 256, 256)
        self.conv_decode2 = expansive_block(256, 128, 128)
        self.conv_decode1 = expansive_block(128, 64, 64)

        self.final_layer = final_block(64, out_channel)

    def forward(self, x):
        # Encode
        encode_block1 = self.conv_encode1(x)
        encode_pool1 = self.conv_pool1(encode_block1)
        encode_block2 = self.conv_encode2(encode_pool1)
        encode_pool2 = self.conv_pool2(encode_block2)
        encode_block3 = self.conv_encode3(encode_pool2)
        encode_pool3 = self.conv_pool3(encode_block3)
        encode_block4 = self.conv_encode4(encode_pool3)
        encode_pool4 = self.conv_pool4(encode_block4)

        # Bottleneck
        bottleneck = self.bottleneck(encode_pool4)

        # Decode
        decode_block4 = self.conv_decode4(encode_block4, bottleneck)
        decode_block3 = self.conv_decode3(encode_block3, decode_block4)
        decode_block2 = self.conv_decode2(encode_block2, decode_block3)
        decode_block1 = self.conv_decode1(encode_block1, decode_block2)

        final_layer = self.final_layer(decode_block1)
        return final_layer

if __name__ == '__main__':
    image = torch.rand((1, 3, 572, 572))
    unet = UNet(in_channel=3, out_channel=2)
    mask = unet(image)
    print(mask.shape)
    
    #输出结果：
    torch.Size([1, 2, 388, 388])

图一：图像裁剪过程演示：
这里演示的是将64x64的特征图裁剪为56x56大小的过程