Unet图像分割网络Pytorch实现

介绍

最近一个月接触了一下Pytorch，个人认为Pytorch相较于Tensorflow来说好用很多。本文的内容是我对Unet论文的总结与提炼，需要提醒的是，Unet原文发布的时候还没有提出BN(Batch Normalization). 所以在本文中我会增加这一个步骤。

如果想要安装Python和Pytorch或者获得进一步的信息可以点击Python ,Pytorch

在图像分割这个大问题上，主要有两个流派：U-shape和dialated Conv。本文介绍的是U-shape网络中最为经典的U-Net。随着骨干网路的进化，很多相应衍生出来的网络大多都是对于Unet进行了改进但是本质上的思路还是没有太多的变化。比如结合DenseNet 和Unet的FCDenseNet, Unet++

---

Unet

Unet是一个为医学图像分割设计的auto-encoder-decoder结构的网络。行业里也把它视作一种FCN(fully connected network)。 它可以分成两个部分，down(encoder) 和 up(decoder)。down的主要结构可以看成conv后面跟maxpool。 up的主要结构是一个upsample后面跟conv。

Unet的核心思想

想要弄清这个问题首先要感性的理解一下卷积的作用。就拿MINIST数据集训练数字识别这个简单的CNN网络为例， 它把一个28*28的图片抽象成一个0-9的向量。卷积可以看成是特征的提取，它可以提取出输入的信息的抽象概念。但是Pool和Conv会损失空间信息。其中，空间信息在pool的过程中损失的更为严重。对于图像分割来说， 空间信息和抽象信息同样重要。既然每一个次pool的时候会严重损失空间信息，也就是说maxpool之间的空间信息多于之后的。于是Unet提出，把down的特征连接到对应的up上。

Unet的结构

其中灰色箭头copy and crop中的copy就是concatenate而crop是为了让两者的长宽一致
左半边就是down path右半边 就是up path。我们来分别介绍这两个部分。

Down Path

图中input image tile就是我们输入的训练数据。除了第一层是两个conv，其他层都可以看成是maxpool后面跟两个conv。在Unet中绝大部分的conv都是两个conv连用的形式存在的，为了方便，我们可以先自定义一个double_conv类。

# 实现double conv
class double_conv(nn.Module):
    ''' Conv => Batch_Norm => ReLU => Conv2d => Batch_Norm => ReLU
    '''
    def __init__(self, in_ch, out_ch):
        super(double_conv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )
        self.conv.apply(self.init_weights)
    
    def forward(self, x):
        x = self.conv(x)
        return x

    @staticmethod
    def init_weights(m):
        if type(m) == nn.Conv2d:
            init.xavier_normal(m.weight)
            init.constant(m.bias,0)

下面我们来实现input conv， 它实际上用一个double_conv也就完成了。

# 实现input conv
class inconv(nn.Module):
    ''' input conv layer
        let input 3 channels image to 64 channels
        The oly difference between `inconv` and `down` is maxpool layer 
    '''
    def __init__(self, in_ch, out_ch):
        super(inconv, self).__init__()
        self.conv = double_conv(in_ch, out_ch)

    def forward