深度学习pytorch框架--Unet医疗影像分割

基于Unet的医疗影像分割

简单复现Unet网络，用来练习pytorch，

  U-net结构（最低分辨率下32x32像素的例子）如下图所示。每个蓝框对应于一个多通道特征图。通道的数量表示在盒子的顶部。X-Y尺寸在盒子的左下角提供。白色方框代表复制的特征图。箭头表示不同的操作。

  其中，蓝/白框表示feature map；蓝色剪头表示3x3 卷积，用于特征提取；灰色箭头表示skip-connection，用于将下采样的浅层特征与上采样层的深层特征融合，提高分割准确性。；红色箭头表示池化，用于池化，绿色箭头表示上采样，用于恢复维度；绿色箭头表示1x1 卷积，用于输出结果。
  unet网络是基于Encoder-Decoder 即编码-解码的设计思想设计的，其中编码部分是由卷积操作和下采样操作组成，文中所用的卷积结构统一为3x3 的卷积核，padding为0， 步长stride为1，但是padding为0导致每次卷积之后造成feature map的H和W都变小了，这会导致边缘信息的丢失，所以我们将padding为1.而在上述的两次卷积之后是一个池化操作用于降低数据维度，需要注意的是在池化之前要保存两次卷积之后的输出用于与上采样层的深层特征融合。

pytorch代码如下

class DownsampleLayer(nn.Module):
    """
    定义下采样忽的网络层
    """
    def __init__(self, in_chanel, out_chanel):
        super(DownsampleLayer, self).__init__()
        self.conv_relu = nn.Sequential(
            nn.Conv2d(in_channels=in_chanel, out_channels=out_chanel, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_chanel),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=out_chanel, out_channels=out_chanel, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_chanel),
            nn.ReLU(inplace=True)
        )
        self.dowmsample = nn.Sequential(
            # 利用步长实现向下池化，利用步长为2卷积避免因数据长度为单数而丢失边缘信息
            # nn.MaxPool2d(2)
            nn.Conv2d(in_channels=out_chanel, out_channels=out_chanel, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(out_chanel),
            nn.ReLU(inplace=True)
        )

  此处的DownsampleLayer层即为下图中的红色方框圈出来的部分

  在红圈部分的的步骤重复四次之后就会进入解码的阶段，其基本层是上图中黄色方框圈出来的部分，他包括两层卷积层和一个反卷积层。其中反卷积层是用来恢复维度的。而下一层的输入是反卷积层输出的结果与下采样层的浅层特征的拼接。

pytorch代码如下

    def __init__(self, in_chanel, out_chanel):
        super(UpsampleLayer, self).__init__()
        self.up_conv_relu = nn.Sequential(
            nn.Conv2d(in_channels=in_chanel, out_channels=out_chanel*2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_chanel * 2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=out_chanel*2, out_channels=out_chanel*2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_chanel * 2),
            nn.ReLU(inplace=True)
        )
        self.upsample = nn.Sequential(
            nn.ConvTranspose2d(in_channels=out_chanel*2, out_channels=out_chanel, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(out_chanel),
            nn.ReLU(inplace=True)
        )

  在黄圈部分重复了四次之后，进行两次卷积核为3x3 的卷积和一次卷积核为1x1的卷积。用于输出结果
代码如下

# 输出
        self.o = nn.Sequential(
            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1, stride=1, padding=1),
        )

结果如下

0_test	0_predict
6_test	6_predict

代码

代码下载地址