本文 GitHub https://github.com/Jack-Cherish/PythonPark 已收錄，有技術乾貨文章，整理的學習資料，一線大廠面試經驗分享等，歡迎 Star 和 完善。
一、前言
本文屬於Pytorch深度學習語義分割系列教程。
該系列文章的內容有：
- Pytorch的基本使用
- 語義分割算法講解
如果不瞭解語義分割原理以及開發環境的搭建，請看該系列教程的上一篇文章《Pytorch深度學習實戰教程（一）：語義分割基礎與環境搭建》。
本文的開發環境採用上一篇文章搭建好的Windows環境，環境情況如下：
開發環境：Windows
開發語言：Python3.7.4
框架版本：Pytorch1.3.0
CUDA：10.2
cuDNN：7.6.0
本文主要講解UNet網絡結構，以及相應代碼的代碼編寫。
PS：文中出現的所有代碼，均可在我的github上下載，歡迎Follow、Star：點擊查看
二、UNet網絡結構
在語義分割領域，基於深度學習的語義分割算法開山之作是FCN（Fully Convolutional Networks for Semantic Segmentation），而UNet是遵循FCN的原理，並進行了相應的改進，使其適應小樣本的簡單分割問題。
UNet論文地址：點擊查看
研究一個深度學習算法，可以先看網絡結構，看懂網絡結構後，再Loss計算方法、訓練方法等。本文主要針對UNet的網絡結構進行講解，其它內容會在後續章節進行說明。
1、網絡結構原理
UNet最早發表在2015的MICCAI會議上，4年多的時間，論文引用量已經達到了9700多次。
UNet成爲了大多做醫療影像語義分割任務的baseline，同時也啓發了大量研究者對於U型網絡結構的研究，發表了一批基於UNet網絡結構的改進方法的論文。
UNet網絡結構，最主要的兩個特點是：U型網絡結構和Skip Connection跳層連接。
UNet是一個對稱的網絡結構，左側爲下采樣，右側爲上採樣。
按照功能可以將左側的一系列下采樣操作稱爲encoder，將右側的一系列上採樣操作稱爲decoder。
Skip Connection中間四條灰色的平行線，Skip Connection就是在上採樣的過程中，融合下采樣過過程中的feature map。
Skip Connection用到的融合的操作也很簡單，就是將feature map的通道進行疊加，俗稱Concat。
Concat操作也很好理解，舉個例子：一本大小爲10cm*10cm，厚度爲3cm的書A，和一本大小爲10cm*10cm，厚度爲4cm的書B。
將書A和書B，邊緣對齊地摞在一起。這樣就得到了，大小爲10cm*10cm厚度爲7cm的一摞書，類似這種：
這種“摞在一起”的操作，就是Concat。
同樣道理，對於feature map，一個大小爲256*256*64的feature map，即feature map的w（寬）爲256，h（高）爲256，c（通道數）爲64。和一個大小爲256*256*32的feature map進行Concat融合，就會得到一個大小爲256*256*96的feature map。
在實際使用中，Concat融合的兩個feature map的大小不一定相同，例如256*256*64的feature map和240*240*32的feature map進行Concat。
這種時候，就有兩種辦法：
第一種：將大256*256*64的feature map進行裁剪，裁剪爲240*240*64的feature map，比如上下左右，各捨棄8 pixel，裁剪後再進行Concat，得到240*240*96的feature map。
第二種：將小240*240*32的feature map進行padding操作，padding爲256*256*32的feature map，比如上下左右，各補8 pixel，padding後再進行Concat，得到256*256*96的feature map。
UNet採用的Concat方案就是第二種，將小的feature map進行padding，padding的方式是補0，一種常規的常量填充。
2、代碼
有些朋友可能對Pytorch不太瞭解，推薦一個快速入門的官方教程。一個小時，你就可以掌握一些基本概念和Pytorch代碼編寫方法。
Pytorch官方基礎：點擊查看
我們將整個UNet網絡拆分爲多個模塊進行講解。
DoubleConv模塊：
先看下連續兩次的卷積操作。
從UNet網絡中可以看出，不管是下采樣過程還是上採樣過程，每一層都會連續進行兩次卷積操作，這種操作在UNet網絡中重複很多次，可以單獨寫一個DoubleConv模塊：
import torch.nn as nn
class DoubleConv(nn.Module):
"""(convolution => [BN] => ReLU) * 2"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.double_conv(x)
解釋下，上述的Pytorch代碼：torch.nn.Sequential是一個時序容器，Modules 會以它們傳入的順序被添加到容器中。比如上述代碼的操作順序：卷積->BN->ReLU->卷積->BN->ReLU。
DoubleConv模塊的in_channels和out_channels可以靈活設定，以便擴展使用。
如上圖所示的網絡，in_channels設爲1，out_channels爲64。
輸入圖片大小爲572*572，經過步長爲1，padding爲0的3*3卷積，得到570*570的feature map，再經過一次卷積得到568*568的feature map。
計算公式：O=(H−F+2×P)/S+1
H爲輸入feature map的大小，O爲輸出feature map的大小，F爲卷積核的大小，P爲padding的大小，S爲步長。
Down模塊：
UNet網絡一共有4次下采樣過程，模塊化代碼如下：
class Down(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels)
)
def forward(self, x):
return self.maxpool_conv(x)
這裏的代碼很簡單，就是一個maxpool池化層，進行下采樣，然後接一個DoubleConv模塊。
至此，UNet網絡的左半部分的下采樣過程的代碼都寫好了，接下來是右半部分的上採樣過程。
Up模塊：
上採樣過程用到的最多的當然就是上採樣了，除了常規的上採樣操作，還有進行特徵的融合。
這塊的代碼實現起來也稍複雜一些：
class Up(nn.Module):
"""Upscaling then double conv"""
def __init__(self, in_channels, out_channels, bilinear=True):
super().__init__()
# if bilinear, use the normal convolutions to reduce the number of channels
if bilinear:
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
else:
self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
self.conv = DoubleConv(in_channels, out_channels)
def forward(self, x1, x2):
x1 = self.up(x1)
# input is CHW
diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
diffX = torch.tensor([x2.size()[3] - x1.size()[3]])
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2])
# if you have padding issues, see
# https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
# https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
代碼複雜一些，我們可以分開來看，首先是__init__初始化函數裏定義的上採樣方法以及卷積採用DoubleConv。上採樣，定義了兩種方法：Upsample和ConvTranspose2d，也就是雙線性插值和反捲積。
雙線性插值很好理解，示意圖：
熟悉雙線性插值的朋友對於這幅圖應該不陌生，簡單地講：已知Q11、Q12、Q21、Q22四個點座標，通過Q11和Q21求R1，再通過Q12和Q22求R2，最後通過R1和R2求P，這個過程就是雙線性插值。
對於一個feature map而言，其實就是在像素點中間補點，補的點的值是多少，是由相鄰像素點的值決定的。
反捲積，顧名思義，就是反着卷積。卷積是讓featuer map越來越小，反捲積就是讓feature map越來越大，示意圖：
下面藍色爲原始圖片，周圍白色的虛線方塊爲padding結果，通常爲0，上面綠色爲卷積後的圖片。
這個示意圖，就是一個從2*2的feature map->4*4的feature map過程。
在forward前向傳播函數中，x1接收的是上採樣的數據，x2接收的是特徵融合的數據。特徵融合方法就是，上文提到的，先對小的feature map進行padding，再進行concat。
OutConv模塊：
用上述的DoubleConv模塊、Down模塊、Up模塊就可以拼出UNet的主體網絡結構了。UNet網絡的輸出需要根據分割數量，整合輸出通道，結果如下圖所示：
操作很簡單，就是channel的變換，上圖展示的是分類爲2的情況（通道爲2）。
雖然這個操作很簡單，也就調用一次，爲了美觀整潔，也封裝一下吧。
class OutConv(nn.Module):
def __init__(self, in_channels, out_channels):
super(OutConv, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
def forward(self, x):
return self.conv(x)
至此，UNet網絡用到的模塊都已經寫好，我們可以將上述的模塊代碼都放到一個unet_parts.py文件裏，然後再創建unet_model.py，根據UNet網絡結構，設置每個模塊的輸入輸出通道個數以及調用順序，編寫如下代碼：
""" Full assembly of the parts to form the complete network """
"""Refer https://github.com/milesial/Pytorch-UNet/blob/master/unet/unet_model.py"""
import torch.nn.functional as F
from unet_parts import *
class UNet(nn.Module):
def __init__(self, n_channels, n_classes, bilinear=False):
super(UNet, self).__init__()
self.n_channels = n_channels
self.n_classes = n_classes
self.bilinear = bilinear
self.inc = DoubleConv(n_channels, 64)
self.down1 = Down(64, 128)
self.down2 = Down(128, 256)
self.down3 = Down(256, 512)
self.down4 = Down(512, 1024)
self.up1 = Up(1024, 512, bilinear)
self.up2 = Up(512, 256, bilinear)
self.up3 = Up(256, 128, bilinear)
self.up4 = Up(128, 64, bilinear)
self.outc = OutConv(64, n_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)
logits = self.outc(x)
return logits
if __name__ == '__main__':
net = UNet(n_channels=3, n_classes=1)
print(net)
使用命令python unet_model.py，如果沒有錯誤，你會得到如下結果：
UNet(
(inc): DoubleConv(
(double_conv): Sequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
(4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
)
(down1): Down(
(maxpool_conv): Sequential(
(0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(1): DoubleConv(
(double_conv): Sequential(
(0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))
(1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
(4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
)
)
)
(down2): Down(
(maxpool_conv): Sequential(
(0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(1): DoubleConv(
(double_conv): Sequential(
(0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1))
(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
(4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
)
)
)
(down3): Down(
(maxpool_conv): Sequential(
(0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(1): DoubleConv(
(double_conv): Sequential(
(0): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1))
(1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1))
(4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
)
)
)
(down4): Down(
(maxpool_conv): Sequential(
(0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(1): DoubleConv(
(double_conv): Sequential(
(0): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1))
(1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))
(4): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
)
)
)
(up1): Up(
(up): ConvTranspose2d(1024, 512, kernel_size=(2, 2), stride=(2, 2))
(conv): DoubleConv(
(double_conv): Sequential(
(0): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1))
(1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1))
(4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
)
)
(up2): Up(
(up): ConvTranspose2d(512, 256, kernel_size=(2, 2), stride=(2, 2))
(conv): DoubleConv(
(double_conv): Sequential(
(0): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1))
(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
(4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
)
)
(up3): Up(
(up): ConvTranspose2d(256, 128, kernel_size=(2, 2), stride=(2, 2))
(conv): DoubleConv(
(double_conv): Sequential(
(0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1))
(1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
(4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
)
)
(up4): Up(
(up): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2))
(conv): DoubleConv(
(double_conv): Sequential(
(0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
(4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): ReLU(inplace=True)
)
)
)
(outc): OutConv(
(conv): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1))
)
)
網絡搭建完成，下一步就是使用網絡進行訓練了，具體實現會在該系列教程的下一篇文章進行講解。
三、小結
- 本文主要講解了UNet網絡結構，並對UNet網絡進行了模塊化梳理。
- 下篇文章講解如何使用UNet網絡，編寫訓練代碼。
點贊再看，養成習慣，微信公衆號搜索【JackCui-AI】關注一個在互聯網摸爬滾打的潛行者
