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的过程看作解码器

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

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在医疗图像分割种表现好

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

4 Unet网络结构

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

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

这里的输入是单通道的原因是因为输入图片是灰度图，而输出是两通道是因为这里是对像素进行二分类（前景和背景），所以输出通道是2
Unet网络解析
整个网络由编码部分（左） 和 解码部分（右）组成，类似于一个大大的U字母，具体介绍如下：

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

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

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大小的过程
Unet网络解析

标签 , ,