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