UNet详细解读（二）pytorch从头开始搭建UNet

Unet代码

网络架构图

• 输入是572x572的，但是输出变成了388x388，这说明经过网络以后，输出的结果和原图不是完全对应的，这在计算loss和输出结果都可以得到体现.

• 蓝色箭头代表3x3的卷积操作，并且步长是1,不进行padding，因此，每个该操作以后，featuremap的大小会减2.

• 红色箭头代表2x2的最大池化操作.如果池化之前特征向量的大小是奇数，那么就会损失一些信息 。输入的大小最好满足一个条件，就是可以让每一层池化操作前的特征向量的大小是偶数，这样就不会损失一些信息，并且crop的时候不会产生误差.

• 绿色箭头代表2x2的反卷积操作.

• 灰色箭头表示复制和剪切操作.

• 输出的最后一层，使用了1x1的卷积层做了分类

• 前半部分也就是图中左边部分的作用是特征提取，后半部分也就是图中的右边部分是上采样，也叫 encoder-deconder结构

两个3X3卷积层

蓝色箭头代表3x3的卷积操作，并且步长是1,不进行padding，因此，每个该操作以后，featuremap的大小会减2.

class DoubleConvolution(nn.Module):

	def __init__(self, in_channels: int, out_channels: int):
    	super().__init__()

		self.first = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
		self.act1 = nn.ReLU()
		self.second = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
		self.act2 = nn.ReLU()
		
	def forward(self, x: torch.Tensor):
		x = self.first(x)
        x = self.act1(x)
        x = self.second(x)
        return self.act2(x)

下采样

红色箭头代表2x2的最大池化操作。

class DownSample(nn.Module):
    def __init__(self):
        super().__init__()
        self.pool = nn.MaxPool2d(2)

    def forward(self,x:torch.Tensor):
        return self.pool(x)

上采样

绿色箭头代表2x2的反卷积操作.

class UpSample(nn.Module):
    def __init__(self,input_channals:int,output_channals:int):
        super().__init__()
        self.up = nn.ConvTranspose2d(input_channals,output_channals,kernel_size=2,stride=2)

    def forward(self,x:torch.Tensor):
        return self.up(x)

裁剪并连接特征图

在扩展路径中的每个步骤，来自收缩路径的对应特征图与当前特征图连接。

contracting_x：将特征图从收缩路径裁剪为当前特征图的大小

class CropAndConcat(nn.Module):

    def forward(self,x:torch.Tensor,contracting_x:torch.Tensor):
        contracting_x = torchvision.transforms.functional.center_crop(contracting_x,[x.shape[2],x.shape[3]])
        x = torch.cat([x,contracting_x],dim=1)
        return x

网络架构代码

class Unet(nn.Module):
    def __init__(self,input_channals:int,output_channals:int):
        super().__init__()

        self.down_conv = nn.ModuleList([DoubleConvolution(i,0) for i,o in [(input_channals,64),(64,128),(128,256),(256,512)]])
        self.down_sample = nn.ModuleList([DownSample() for _ in range(4)])
        self.middel_conv = DoubleConvolution(512,1024)
        self.up_sample = nn.ModuleList([UpSample(i,o) for i,o in [(1024,512),(512,256),(256,128),(128,64)]])
        self.up_conv = nn.ModuleList([DoubleConvolution(i,o) for i,o in [(1024,512),(512,256),(256,128),(128,64)]])
        self.concat = nn.ModuleList(CropAndConcat() for _ in range(4))
        self.final_conv = nn.Conv2d(64,output_channals,kernel_size=1)

    def forward(self,x:torch.Tensor):
        pass_through = []

        for i in range(len(self.down_conv)):
            x = self.down_conv[i](x)
            pass_through.append(x)
            x = self.down_sample[i](x)

        x = self.middel_conv(x)

        for i in range(len(self.up_conv)):
            x = self.up_sample[i](x)
            x = self.concat[i](x,pass_through.pop())
            x = self.up_conv[i](x)

        x = self.final_conv(x)

        return x