UNet詳細解讀（二）pytorch從頭開始搭建UNet

pytorch從頭開始搭建UNet

• Unet代碼
• • 網絡架構圖
  • 兩個3X3卷積層
  • 下采樣
  • 上采樣
  • 裁剪并連接配接特征圖
  • 網絡架構代碼

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