如何分析Pytorch中UNet網(wǎng)絡(luò)結(jié)構(gòu)以及代碼編寫

這篇文章給大家介紹如何分析Pytorch中UNet網(wǎng)絡(luò)結(jié)構(gòu)以及代碼編寫，內(nèi)容非常詳細(xì)，感興趣的小伙伴們可以參考借鑒，希望對(duì)大家能有所幫助。

一、前言

Windows環(huán)境開(kāi)發(fā)，環(huán)境情況如下：

開(kāi)發(fā)環(huán)境：Windows

開(kāi)發(fā)語(yǔ)言：Python3.7.4

框架版本：Pytorch2.3.0

CUDA：10.2

cuDNN：7.6.0

主要講解UNet網(wǎng)絡(luò)結(jié)構(gòu)，以及相應(yīng)代碼的代碼編寫。

PS：文中出現(xiàn)的所有代碼，均可在我的github上下載，歡迎Follow、Star：點(diǎn)擊查看

二、UNet網(wǎng)絡(luò)結(jié)構(gòu)

在語(yǔ)義分割領(lǐng)域，基于深度學(xué)習(xí)的語(yǔ)義分割算法開(kāi)山之作是FCN（Fully Convolutional Networks for Semantic Segmentation），而UNet是遵循FCN的原理，并進(jìn)行了相應(yīng)的改進(jìn)，使其適應(yīng)小樣本的簡(jiǎn)單分割問(wèn)題。

UNet論文地址：點(diǎn)擊查看

研究一個(gè)深度學(xué)習(xí)算法，可以先看網(wǎng)絡(luò)結(jié)構(gòu)，看懂網(wǎng)絡(luò)結(jié)構(gòu)后，再Loss計(jì)算方法、訓(xùn)練方法等。本文主要針對(duì)UNet的網(wǎng)絡(luò)結(jié)構(gòu)進(jìn)行講解，其它內(nèi)容會(huì)在后續(xù)章節(jié)進(jìn)行說(shuō)明。

1、網(wǎng)絡(luò)結(jié)構(gòu)原理

UNet最早發(fā)表在2015的MICCAI會(huì)議上，4年多的時(shí)間，論文引用量已經(jīng)達(dá)到了9700多次。

UNet成為了大多做醫(yī)療影像語(yǔ)義分割任務(wù)的baseline，同時(shí)也啟發(fā)了大量研究者對(duì)于U型網(wǎng)絡(luò)結(jié)構(gòu)的研究，發(fā)表了一批基于UNet網(wǎng)絡(luò)結(jié)構(gòu)的改進(jìn)方法的論文。

UNet網(wǎng)絡(luò)結(jié)構(gòu)，最主要的兩個(gè)特點(diǎn)是：U型網(wǎng)絡(luò)結(jié)構(gòu)和Skip Connection跳層連接。

這種“摞在一起”的操作，就是Concat。

同樣道理，對(duì)于feature map，一個(gè)大小為256*256*64的feature map，即feature map的w（寬）為256，h（高）為256，c（通道數(shù)）為64。和一個(gè)大小為256*256*32的feature map進(jìn)行Concat融合，就會(huì)得到一個(gè)大小為256*256*96的feature map。

在實(shí)際使用中，Concat融合的兩個(gè)feature map的大小不一定相同，例如256*256*64的feature map和240*240*32的feature map進(jìn)行Concat。

這種時(shí)候，就有兩種辦法：

第一種：將大256*256*64的feature map進(jìn)行裁剪，裁剪為240*240*64的feature map，比如上下左右，各舍棄8 pixel，裁剪后再進(jìn)行Concat，得到240*240*96的feature map。

第二種：將小240*240*32的feature map進(jìn)行padding操作，padding為256*256*32的feature map，比如上下左右，各補(bǔ)8 pixel，padding后再進(jìn)行Concat，得到256*256*96的feature map。

UNet采用的Concat方案就是第二種，將小的feature map進(jìn)行padding，padding的方式是補(bǔ)0，一種常規(guī)的常量填充。

2、代碼

有些朋友可能對(duì)Pytorch不太了解，推薦一個(gè)快速入門的官方教程。一個(gè)小時(shí)，你就可以掌握一些基本概念和Pytorch代碼編寫方法。

Pytorch官方基礎(chǔ)：點(diǎn)擊查看

我們將整個(gè)UNet網(wǎng)絡(luò)拆分為多個(gè)模塊進(jìn)行講解。

DoubleConv模塊：

先看下連續(xù)兩次的卷積操作。

從UNet網(wǎng)絡(luò)中可以看出，不管是下采樣過(guò)程還是上采樣過(guò)程，每一層都會(huì)連續(xù)進(jìn)行兩次卷積操作，這種操作在UNet網(wǎng)絡(luò)中重復(fù)很多次，可以單獨(dú)寫一個(gè)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是一個(gè)時(shí)序容器，Modules 會(huì)以它們傳入的順序被添加到容器中。比如上述代碼的操作順序：卷積->BN->ReLU->卷積->BN->ReLU。

DoubleConv模塊的in_channels和out_channels可以靈活設(shè)定，以便擴(kuò)展使用。

如上圖所示的網(wǎng)絡(luò)，in_channels設(shè)為1，out_channels為64。

輸入圖片大小為572*572，經(jīng)過(guò)步長(zhǎng)為1，padding為0的3*3卷積，得到570*570的feature map，再經(jīng)過(guò)一次卷積得到568*568的feature map。

計(jì)算公式：O=(H?F+2×P)/S+1

H為輸入feature map的大小，O為輸出feature map的大小，F(xiàn)為卷積核的大小，P為padding的大小，S為步長(zhǎng)。

Down模塊：

UNet網(wǎng)絡(luò)一共有4次下采樣過(guò)程，模塊化代碼如下：

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)

這里的代碼很簡(jiǎn)單，就是一個(gè)maxpool池化層，進(jìn)行下采樣，然后接一個(gè)DoubleConv模塊。

至此，UNet網(wǎng)絡(luò)的左半部分的下采樣過(guò)程的代碼都寫好了，接下來(lái)是右半部分的上采樣過(guò)程。

Up模塊：

上采樣過(guò)程用到的最多的當(dāng)然就是上采樣了，除了常規(guī)的上采樣操作，還有進(jìn)行特征的融合。

這塊的代碼實(shí)現(xiàn)起來(lái)也稍復(fù)雜一些：

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)

代碼復(fù)雜一些，我們可以分開(kāi)來(lái)看，首先是__init__初始化函數(shù)里定義的上采樣方法以及卷積采用DoubleConv。上采樣，定義了兩種方法：Upsample和ConvTranspose2d，也就是雙線性插值和反卷積。

雙線性插值很好理解，示意圖：

熟悉雙線性插值的朋友對(duì)于這幅圖應(yīng)該不陌生，簡(jiǎn)單地講：已知Q11、Q12、Q21、Q22四個(gè)點(diǎn)坐標(biāo)，通過(guò)Q11和Q21求R1，再通過(guò)Q12和Q22求R2，最后通過(guò)R1和R2求P，這個(gè)過(guò)程就是雙線性插值。

對(duì)于一個(gè)feature map而言，其實(shí)就是在像素點(diǎn)中間補(bǔ)點(diǎn)，補(bǔ)的點(diǎn)的值是多少，是由相鄰像素點(diǎn)的值決定的。

反卷積，顧名思義，就是反著卷積。卷積是讓featuer map越來(lái)越小，反卷積就是讓feature map越來(lái)越大，示意圖：

下面藍(lán)色為原始圖片，周圍白色的虛線方塊為padding結(jié)果，通常為0，上面綠色為卷積后的圖片。

這個(gè)示意圖，就是一個(gè)從2*2的feature map->4*4的feature map過(guò)程。

在forward前向傳播函數(shù)中，x1接收的是上采樣的數(shù)據(jù)，x2接收的是特征融合的數(shù)據(jù)。特征融合方法就是，上文提到的，先對(duì)小的feature map進(jìn)行padding，再進(jìn)行concat。

OutConv模塊：

用上述的DoubleConv模塊、Down模塊、Up模塊就可以拼出UNet的主體網(wǎng)絡(luò)結(jié)構(gòu)了。UNet網(wǎng)絡(luò)的輸出需要根據(jù)分割數(shù)量，整合輸出通道，結(jié)果如下圖所示：

操作很簡(jiǎn)單，就是channel的變換，上圖展示的是分類為2的情況（通道為2）。

雖然這個(gè)操作很簡(jiǎn)單，也就調(diào)用一次，為了美觀整潔，也封裝一下吧。

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網(wǎng)絡(luò)用到的模塊都已經(jīng)寫好，我們可以將上述的模塊代碼都放到一個(gè)unet_parts.py文件里，然后再創(chuàng)建unet_model.py，根據(jù)UNet網(wǎng)絡(luò)結(jié)構(gòu)，設(shè)置每個(gè)模塊的輸入輸出通道個(gè)數(shù)以及調(diào)用順序，編寫如下代碼：

""" 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，如果沒(méi)有錯(cuò)誤，你會(huì)得到如下結(jié)果：

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))
  )
)

關(guān)于如何分析Pytorch中UNet網(wǎng)絡(luò)結(jié)構(gòu)以及代碼編寫就分享到這里了，希望以上內(nèi)容可以對(duì)大家有一定的幫助，可以學(xué)到更多知識(shí)。如果覺(jué)得文章不錯(cuò)，可以把它分享出去讓更多的人看到。