手撸网络结构—UNet

用于分类任务的卷积神经网络输出一般都是一个单一的类标签，但是在很多图像视觉任务中往往要求输出信息包含位置信息，比如给图像的每个像素点赋予类别信息。

这样就需要使用到fully convolutional network 相较于FCN，UNet不只是简单的将卷积网络的结果直接进行转置卷积，而是结合了前面卷积层的中间计算结果，这样既利用了卷积部分的位置信息又利用了最后的分类信息。

UNet网络结构

UNet分为两个阶段：downsample stage 和 upsample stage。

图中的每个蓝色框对应了每层的feature map，上面数字表示当前的通道数。左下角的数字表示feature map的尺寸，这里输入为572x572通道数为1。白色的框表示将对应的卷积层结果进行crop后的结果，白色框与蓝色框进行组合，作为后续网络的输入。

因为在前面卷积的过程中没有添加padding，因此会存在feature map的尺寸变化，所以在上采样过程中，需要对每个阶段的卷积计算结果进行crop操作然后再与转置卷积的结果进行concat（通道维度的相加）。

def concat(tensor1,tensor2):
        # concat 2 tensor by the channel axes
        tensor1,tensor2 = (tensor1,tensor2) if tensor1.size()[3]>=tensor2.size()[3] else (tensor2,tensor1)
        crop_val = int((tensor1.size()[3]-tensor2.size()[3])/2)
        tensor1 = tensor1[:, :, crop_val:tensor1.size()[3]-crop_val
                      , crop_val:tensor1.size()[3]-crop_val]
        return torch.cat((tensor1,tensor2),1)

完整代码：

import torch
import torch.nn as nn
from torchsummary import summary


def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride)


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution without padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride)

def up_conv2x2(in_planes,out_planes):
    return nn.ConvTranspose2d(in_planes,out_planes,kernel_size=2,stride=2)

def max_pool2x2():
    return nn.MaxPool2d(kernel_size=2,stride=2)


class UNet(nn.Module):
    def __init__(self,class_num=1000):
        super(UNet, self).__init__()
        # downsample stage
        self.conv_1 = nn.Sequential(conv3x3(1,64),conv3x3(64,64))
        self.conv_2 = nn.Sequential(conv3x3(64,128),conv3x3(128,128))
        self.conv_3 = nn.Sequential(conv3x3(128,256),conv3x3(256,256))
        self.conv_4 = nn.Sequential(conv3x3(256,512),conv3x3(512,512))
        self.conv_5 = nn.Sequential(conv3x3(512,1024),conv3x3(1024,1024))
        self.maxpool = max_pool2x2()
        
        # upsample stage
        # up_conv_4 corresponds conv_4
        self.up_conv_4 = nn.Sequential(up_conv2x2(1024,512))
        # conv the cat(stage_4,up_conv_4) from 1024 to 512
        self.conv_6 = nn.Sequential(conv3x3(1024,512),conv3x3(512,512))
        # up_conv_3 corresponds conv_3
        self.up_conv_3 = nn.Sequential(up_conv2x2(512,256))
        # conv the cat(stage_3,up_conv_3) from 512 to 256
        self.conv_7 = nn.Sequential(conv3x3(512,256),conv3x3(256,256))
        # up_conv_2 corresponds conv_2
        self.up_conv_2 = nn.Sequential(up_conv2x2(256,128))
        # conv the cat(stage_2,up_conv_2) from 256 to 128
        self.conv_8 = nn.Sequential(conv3x3(256,128),conv3x3(128,128))
        # up_conv_1 corresponds conv_1
        self.up_conv_1 = nn.Sequential(up_conv2x2(128,64))
        # conv the cat(stage_1,up_conv_1) from 128 to 64
        self.conv_9 = nn.Sequential(conv3x3(128,64),conv3x3(64,64))
        # output
        self.result = conv1x1(64,2)
    
    def _concat(self,tensor1,tensor2):
        # concat 2 tensor by the channel axes
        tensor1,tensor2 = (tensor1,tensor2) if tensor1.size()[3]>=tensor2.size()[3] else (tensor2,tensor1)
        crop_val = int((tensor1.size()[3]-tensor2.size()[3])/2)
        tensor1 = tensor1[:, :, crop_val:tensor1.size()[3]-crop_val
                      , crop_val:tensor1.size()[3]-crop_val]
        return torch.cat((tensor1,tensor2),1)

    def forward(self,x):
        # get 4 stage conv output
        stage_1 = self.conv_1(x)
        stage_2 = self.conv_2(self.maxpool(stage_1))
        stage_3 = self.conv_3(self.maxpool(stage_2))
        stage_4 = self.conv_4(self.maxpool(stage_3))

        # get up_conv_4 and concat with stage_4
        up_in_4 = self.conv_5(self.maxpool(stage_4))
        up_stage_4 = self.up_conv_4(up_in_4)
        up_stage_4 = self._concat(stage_4,up_stage_4)
        # get up_conv_3 and concat with stage_3
        up_in_3 = self.conv_6(up_stage_4)
        up_stage_3 = self.up_conv_3(up_in_3)
        up_stage_3 = self._concat(stage_3,up_stage_3)
        # get up_conv_2 and concat with stage_2
        up_in_2 = self.conv_7(up_stage_3)
        up_stage_2 = self.up_conv_2(up_in_2)
        up_stage_2 = self._concat(stage_2,up_stage_2)
        # get up_conv_1 and concat with stage_1
        up_in_1 = self.conv_8(up_stage_2)
        up_stage_1 = self.up_conv_1(up_in_1)
        up_stage_1 = self._concat(stage_1,up_stage_1)

        # last conv to channel 2
        out = self.conv_9(up_stage_1)
        # result
        out = self.result(out)
        return out


if __name__ == '__main__':
    ut = UNet(12)
    summary(ut,(1,572,572))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 570, 570]             640
            Conv2d-2         [-1, 64, 568, 568]          36,928
         MaxPool2d-3         [-1, 64, 284, 284]               0
            Conv2d-4        [-1, 128, 282, 282]          73,856
            Conv2d-5        [-1, 128, 280, 280]         147,584
         MaxPool2d-6        [-1, 128, 140, 140]               0
            Conv2d-7        [-1, 256, 138, 138]         295,168
            Conv2d-8        [-1, 256, 136, 136]         590,080
         MaxPool2d-9          [-1, 256, 68, 68]               0
           Conv2d-10          [-1, 512, 66, 66]       1,180,160
           Conv2d-11          [-1, 512, 64, 64]       2,359,808
        MaxPool2d-12          [-1, 512, 32, 32]               0
           Conv2d-13         [-1, 1024, 30, 30]       4,719,616
           Conv2d-14         [-1, 1024, 28, 28]       9,438,208
  ConvTranspose2d-15          [-1, 512, 56, 56]       2,097,664
           Conv2d-16          [-1, 512, 54, 54]       4,719,104
           Conv2d-17          [-1, 512, 52, 52]       2,359,808
  ConvTranspose2d-18        [-1, 256, 104, 104]         524,544
           Conv2d-19        [-1, 256, 102, 102]       1,179,904
           Conv2d-20        [-1, 256, 100, 100]         590,080
  ConvTranspose2d-21        [-1, 128, 200, 200]         131,200
           Conv2d-22        [-1, 128, 198, 198]         295,040
           Conv2d-23        [-1, 128, 196, 196]         147,584
  ConvTranspose2d-24         [-1, 64, 392, 392]          32,832
           Conv2d-25         [-1, 64, 390, 390]          73,792
           Conv2d-26         [-1, 64, 388, 388]          36,928
           Conv2d-27          [-1, 2, 388, 388]             130
================================================================
Total params: 31,030,658
Trainable params: 31,030,658
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 1.25
Forward/backward pass size (MB): 1096.59
Params size (MB): 118.37
Estimated Total Size (MB): 1216.21
----------------------------------------------------------------
ckward pass size (MB): 1096.59
Params size (MB): 118.37
Estimated Total Size (MB): 1216.21
----------------------------------------------------------------