U-Net是一種用于圖像分割的深度學習模型,特別適合醫學影像和其他需要分割細節的任務。如圖:
Unet論文原文
為什么叫U-Net?
U-Net的結構像字母“U”,所以得名。它的結構由兩個主要部分組成:
下采樣(編碼器):圖像逐漸被縮小并且提取特征。
上采樣(解碼器):逐漸恢復圖像的尺寸,并通過“跳躍連接”將高分辨率的特征與低分辨率的特征結合,以保持細節。
網絡結構
U-Net通常包括以下幾部分:
(1)下采樣(編碼器):
這一部分通過卷積操作逐步降低圖像的分辨率,同時提取越來越抽象的特征。每個卷積層后面跟著一個池化層(比如最大池化),池化層會減少圖像的尺寸,幫助網絡捕捉全局信息。
(2)瓶頸層:
當圖像經過幾次下采樣后,網絡會到達最小尺寸的圖像,這一部分的特征信息會被壓縮成最具代表性的特征。
(3)上采樣(解碼器):
上采樣過程會將圖像逐漸恢復到原始大小,通過轉置卷積(反卷積)或者插值操作來擴大圖像的尺寸。
(4)跳躍連接:
這里的“跳躍連接”是U-Net的一大特點。在上采樣階段,網絡會將下采樣階段的特征圖直接傳遞到解碼器中,這樣可以幫助網絡更好地恢復圖像的細節,防止圖像中的細微特征丟失。
Unet的特點
- UNet也是一個編解碼模型,它將各級分辨率的卷積特征圖和反卷積特征圖進行了融合。
- Unet輸入和輸出大小是同樣的維度。
- 如圖
網絡的右側虛線框部分叫作擴展路徑(expansive path)。
它同樣由4個網絡塊組成,每個網絡塊開始之前通過反卷積將特征圖的尺寸乘2,同時將其通道數減半(最后一層略有不同)?,然后和左側對稱的壓縮路徑的特征圖合并。由于左側壓縮路徑和右側擴展路徑的特征圖的尺寸不一樣,U-Net是通過將壓縮路徑的特征圖裁剪到和擴展路徑的特征圖相同尺寸進行歸一化的。
擴展路徑的卷積操作依舊使用的是有效卷積操作,最終得到的特征圖的尺寸是388×388。由于該任務是一個二分類任務,因此網絡有兩個輸出特征圖。
在U-Net的網絡結構中通過裁剪的方式將左側的特征圖調整到和右側特征圖相同的大小(圖左側部分中的虛線部分)
U-Net論文中沒有全連接層(Linear層),且全程使用valid來進行卷積,這樣可以保證分割的結果都是基于沒有缺失的上下文特征得到的,因此輸入輸出的圖像尺寸不一樣。
需要注意的是,論文中的輸入圖像的尺寸為 572 像素×572 像素,而輸出圖像的尺寸是388 像素×388像素大小,輸入與輸出并不一致;整個網絡在很多情況下,輸出比輸入寬高減少了 2,這主要是因為卷積層采用了“valid”的邊緣填充(padding)方式,其中 stride=1,因此,輸出的特征圖會逐漸減小。
評價指標
Dice Loss是由Dice系數而得名的,Dice系數是一種用于評估兩個樣本相似性的度量函數,其值越大意味著這兩個樣本越相似,Dice系數的數學表達式如下:
式中,X表示真實分割圖像的像素標簽,Y表示模型預測分割圖像的像素類別,|X∩Y|為預測圖像的像素與真實標簽圖像的像素之間的點乘,并將點乘結果相加,|X|和|Y|分別為它們各自對應圖像中的像素相加。
代碼舉例
class UNet(nn.Module):def __init__(self, in_channels, num_classes):super(UNet, self).__init__()self.inc = InConv(in_channels, 64)self.down1 = downsample_block(64, 128)self.down2 = downsample_block(128, 256)self.down3 = downsample_block(256, 512)self.down4 = downsample_block(512, 512)self.up1 = upsample_block(1024, 256)self.up2 = upsample_block(512, 128)self.up3 = upsample_block(256, 64)self.up4 = upsample_block(128, 64)self.outc = OutConv(64, num_classes)def forward(self, x):x1 = self.inc(x)x2 = self.down1(x1)x3 = self.down2(x2)x4 = self.down3(x3)x5 = self.down4(x4)x = self.up1(x5, x4)x = self.up2(x, x3)x = self.up3(x, x2)x = self.up4(x, x1)x = self.outc(x)return x
in_channels 是輸入圖像的通道數,num_classes 是輸出類別數(通常為1)
U-Net的前向傳播過程
編碼器路徑:
輸入圖像首先通過 InConv 模塊,生成64通道的特征圖 x1。
然后依次通過四個下采樣模塊 down1、down2、down3、down4,分別生成128、256、512、512通道的特征圖 x2、x3、x4、x5。
解碼器路徑:
解碼器路徑從 x5 開始,通過 up1 模塊與 x4 進行拼接,生成256通道的特征圖。
然后依次通過 up2、up3、up4 模塊,分別與 x3、x2、x1 進行拼接,生成128、64、64通道的特征圖。
最后通過 OutConv 模塊生成最終的分割結果。
其中,函數定義如下: