如何深度解析Pytorch中的UNet模型

這篇文章給大家介紹如何深度解析Pytorch中的UNet模型，內容非常詳細，感興趣的小伙伴們可以參考借鑒，希望對大家能有所幫助。

一、項目背景

深度學習算法，無非就是我們解決一個問題的方法。選擇什么樣的網絡去訓練，進行什么樣的預處理，采用什么Loss和優化方法，都是根據具體的任務而定的。

所以，讓我們先看一下今天的任務。

沒錯，就是 UNet 論文中的經典任務：醫學圖像分割。

選擇它作為今天的任務，就是因為簡單，好上手。

簡單描述一個這個任務：如動圖所示，給一張細胞結構圖，我們要把每個細胞互相分割開來。

這個訓練數據只有30張，分辨率為512x512，這些圖片是果蠅的電鏡圖。

好了，任務介紹完畢，開始準備訓練模型。

二、UNet訓練

想要訓練一個深度學習模型，可以簡單分為三個步驟：

• 數據加載：數據怎么加載，標簽怎么定義，用什么數據增強方法，都是這一步進行。

• 模型選擇：模型我們已經準備好了，就是該系列上篇文章講到的 UNet 網絡。

• 算法選擇：算法選擇也就是我們選什么 loss ，用什么優化算法。

每個步驟說的比較籠統，我們結合今天的醫學圖像分割任務，展開說明。

1、數據加載

這一步，可以做很多事情，說白了，無非就是圖片怎么加載，標簽怎么定義，為了增加算法的魯棒性或者增加數據集，可以做一些數據增強的操作。

既然是處理數據，那么我們先看下數據都是什么樣的，再決定怎么處理。

數據已經備好，都在這里（Github）：點擊查看

如果 Github 下載速度慢，可以使用文末的百度鏈接下載數據集。

數據分為訓練集和測試集，各30張，訓練集有標簽，測試集沒有標簽。

數據加載要做哪些處理，是根據任務和數據集而決定的，對于我們的分割任務，不用做太多處理，但由于數據量很少，僅30張，我們可以使用一些數據增強方法，來擴大我們的數據集。

Pytorch 給我們提供了一個方法，方便我們加載數據，我們可以使用這個框架，去加載我們的數據。看下偽代碼：

# ================================================================== #
#                Input pipeline for custom dataset                 #
# ================================================================== #

# You should build your custom dataset as below.
class CustomDataset(torch.utils.data.Dataset):
    def __init__(self):
        # TODO
        # 1. Initialize file paths or a list of file names. 
        pass
    def __getitem__(self, index):
        # TODO
        # 1. Read one data from file (e.g. using numpy.fromfile, PIL.Image.open).
        # 2. Preprocess the data (e.g. torchvision.Transform).
        # 3. Return a data pair (e.g. image and label).
        pass
    def __len__(self):
        # You should change 0 to the total size of your dataset.
        return 0 

# You can then use the prebuilt data loader. 
custom_dataset = CustomDataset()
train_loader = torch.utils.data.DataLoader(dataset=custom_dataset,
                                           batch_size=64, 
                                           shuffle=True)

這是一個標準的模板，我們就使用這個模板，來加載數據，定義標簽，以及進行數據增強。

創建一個dataset.py文件，編寫代碼如下：

import torch
import cv2
import os
import glob
from torch.utils.data import Dataset
import random

class ISBI_Loader(Dataset):
    def __init__(self, data_path):
        # 初始化函數，讀取所有data_path下的圖片
        self.data_path = data_path
        self.imgs_path = glob.glob(os.path.join(data_path, 'image/*.png'))

    def augment(self, image, flipCode):
        # 使用cv2.flip進行數據增強，filpCode為1水平翻轉，0垂直翻轉，-1水平+垂直翻轉
        flip = cv2.flip(image, flipCode)
        return flip
        
    def __getitem__(self, index):
        # 根據index讀取圖片
        image_path = self.imgs_path[index]
        # 根據image_path生成label_path
        label_path = image_path.replace('image', 'label')
        # 讀取訓練圖片和標簽圖片
        image = cv2.imread(image_path)
        label = cv2.imread(label_path)
        # 將數據轉為單通道的圖片
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        label = cv2.cvtColor(label, cv2.COLOR_BGR2GRAY)
        image = image.reshape(1, image.shape[0], image.shape[1])
        label = label.reshape(1, label.shape[0], label.shape[1])
        # 處理標簽，將像素值為255的改為1
        if label.max() > 1:
            label = label / 255
        # 隨機進行數據增強，為2時不做處理
        flipCode = random.choice([-1, 0, 1, 2])
        if flipCode != 2:
            image = self.augment(image, flipCode)
            label = self.augment(label, flipCode)
        return image, label

    def __len__(self):
        # 返回訓練集大小
        return len(self.imgs_path)

    
if __name__ == "__main__":
    isbi_dataset = ISBI_Loader("data/train/")
    print("數據個數：", len(isbi_dataset))
    train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
                                               batch_size=2, 
                                               shuffle=True)
    for image, label in train_loader:
        print(image.shape)

運行代碼，你可以看到如下結果：

解釋一下代碼：

__init__函數是這個類的初始化函數，根據指定的圖片路徑，讀取所有圖片數據，存放到self.imgs_path列表中。

__len__函數可以返回數據的多少，這個類實例化后，通過len()函數調用。

__getitem__函數是數據獲取函數，在這個函數里你可以寫數據怎么讀，怎么處理，并且可以一些數據預處理、數據增強都可以在這里進行。我這里的處理很簡單，只是將圖片讀取，并處理成單通道圖片。同時，因為 label 的圖片像素點是0和255，因此需要除以255，變成0和1。同時，隨機進行了數據增強。

augment函數是定義的數據增強函數，怎么處理都行，我這里只是進行了簡單的旋轉操作。

在這個類中，你不用進行一些打亂數據集的操作，也不用管怎么按照 batchsize 讀取數據。因為實例化這個類后，我們可以用 torch.utils.data.DataLoader 方法指定 batchsize 的大小，決定是否打亂數據。

Pytorch 提供給給我們的 DataLoader 很強大，我們甚至可以指定使用多少個進程加載數據，數據是否加載到 CUDA 內存中等高級用法，本文不涉及，就不再展開講解了。

2、模型選擇

模型我們已經選擇完了，就用上篇文章《Pytorch深度學習實戰教程（二）：UNet語義分割網絡》講解的 UNet 網絡結構。

但是我們需要對網絡進行微調，完全按照論文的結構，模型輸出的尺寸會稍微小于圖片輸入的尺寸，如果使用論文的網絡結構需要在結果輸出后，做一個 resize 操作。為了省去這一步，我們可以修改網絡，使網絡的輸出尺寸正好等于圖片的輸入尺寸。

創建unet_parts.py文件，編寫如下代碼：

""" Parts of the U-Net model """
"""https://github.com/milesial/Pytorch-UNet/blob/master/unet/unet_parts.py"""

import torch
import torch.nn as nn
import torch.nn.functional as F

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=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

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)

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 // 2, 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])

        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


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_model.py文件，編寫如下代碼：

""" 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=True):
        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, 512)
        self.up1 = Up(1024, 256, bilinear)
        self.up2 = Up(512, 128, bilinear)
        self.up3 = Up(256, 64, 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)

這樣調整過后，網絡的輸出尺寸就與圖片的輸入尺寸相同了。

3、算法選擇

選擇什么 Loss 很重要，Loss 選擇的好壞，都會影響算法擬合數據的效果。

選擇什么 Loss 也是根據任務而決定的。我們今天的任務，只需要分割出細胞邊緣，也就是一個很簡單的二分類任務，所以我們可以使用 BCEWithLogitsLoss。

啥是 BCEWithLogitsLoss？BCEWithLogitsLoss 是 Pytorch 提供的用來計算二分類交叉熵的函數。

它的公式是：

看過我機器學習系列教程的朋友，對這個公式一定不陌生，它就是 Logistic 回歸的損失函數。它利用的是 Sigmoid 函數閾值在[0,1]這個特性來進行分類的。

具體的公式推導，可以看我的機器學習系列教程《機器學習實戰教程（六）：Logistic回歸基礎篇之梯度上升算法》，這里就不再累述。

目標函數，也就是 Loss 確定好了，怎么去優化這個目標呢？

最簡單的方法就是，我們耳熟能詳的梯度下降算法，逐漸逼近局部的極值。

但是這種簡單的優化算法，求解速度慢，也就是想找到最優解，費勁兒。

各種優化算法，本質上其實都是梯度下降，例如最常規的 SGD，就是基于梯度下降改進的隨機梯度下降算法，Momentum 就是引入了動量的 SGD，以指數衰減的形式累計歷史梯度。

除了這些最基本的優化算法，還有自適應參數的優化算法。這類算法最大的特點就是，每個參數有不同的學習率，在整個學習過程中自動適應這些學習率，從而達到更好的收斂效果。

本文就是選擇了一種自適應的優化算法 RMSProp。

由于篇幅有限，這里就不再擴展，講解這個優化算法單寫一篇都不夠，要弄懂 RMSProp，你得先知道什么是 AdaGrad，因為 RMSProp 是基于 AdaGrad 的改進。

比 RMSProp 更高級的優化算法也有，比如大名鼎鼎的 Adam，它可以看做是修正后的Momentum+RMSProp 算法。

總之，對于初學者，你只要知道 RMSProp 是一種自適應的優化算法，比較高級就行了。

下面，我們就可以開始寫訓練UNet的代碼了，創建 train.py 編寫如下代碼：

from model.unet_model import UNet
from utils.dataset import ISBI_Loader
from torch import optim
import torch.nn as nn
import torch

def train_net(net, device, data_path, epochs=40, batch_size=1, lr=0.00001):
    # 加載訓練集
    isbi_dataset = ISBI_Loader(data_path)
    train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
                                               batch_size=batch_size, 
                                               shuffle=True)
    # 定義RMSprop算法
    optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.9)
    # 定義Loss算法
    criterion = nn.BCEWithLogitsLoss()
    # best_loss統計，初始化為正無窮
    best_loss = float('inf')
    # 訓練epochs次
    for epoch in range(epochs):
        # 訓練模式
        net.train()
        # 按照batch_size開始訓練
        for image, label in train_loader:
            optimizer.zero_grad()
            # 將數據拷貝到device中
            image = image.to(device=device, dtype=torch.float32)
            label = label.to(device=device, dtype=torch.float32)
            # 使用網絡參數，輸出預測結果
            pred = net(image)
            # 計算loss
            loss = criterion(pred, label)
            print('Loss/train', loss.item())
            # 保存loss值最小的網絡參數
            if loss < best_loss:
                best_loss = loss
                torch.save(net.state_dict(), 'best_model.pth')
            # 更新參數
            loss.backward()
            optimizer.step()

if __name__ == "__main__":
    # 選擇設備，有cuda用cuda，沒有就用cpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加載網絡，圖片單通道1，分類為1。
    net = UNet(n_channels=1, n_classes=1)
    # 將網絡拷貝到deivce中
    net.to(device=device)
    # 指定訓練集地址，開始訓練
    data_path = "data/train/"
    train_net(net, device, data_path)

為了讓工程更加清晰簡潔，我們創建一個 model 文件夾，里面放模型相關的代碼，也就是我們的網絡結構代碼，unet_parts.py 和 unet_model.py。

創建一個 utils 文件夾，里面放工具相關的代碼，比如數據加載工具dataset.py。

這種模塊化的管理，大大提高了代碼的可維護性。

train.py 放在工程根目錄即可，簡單解釋下代碼。

由于數據就30張，我們就不分訓練集和驗證集了，我們保存訓練集 loss 值最低的網絡參數作為最佳模型參數。

如果都沒有問題，你可以看到 loss 正在逐漸收斂。

三、預測

模型訓練好了，我們可以用它在測試集上看下效果。

在工程根目錄創建 predict.py 文件，編寫如下代碼：

import glob
import numpy as np
import torch
import os
import cv2
from model.unet_model import UNet

if __name__ == "__main__":
    # 選擇設備，有cuda用cuda，沒有就用cpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加載網絡，圖片單通道，分類為1。
    net = UNet(n_channels=1, n_classes=1)
    # 將網絡拷貝到deivce中
    net.to(device=device)
    # 加載模型參數
    net.load_state_dict(torch.load('best_model.pth', map_location=device))
    # 測試模式
    net.eval()
    # 讀取所有圖片路徑
    tests_path = glob.glob('data/test/*.png')
    # 遍歷所有圖片
    for test_path in tests_path:
        # 保存結果地址
        save_res_path = test_path.split('.')[0] + '_res.png'
        # 讀取圖片
        img = cv2.imread(test_path)
        # 轉為灰度圖
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # 轉為batch為1，通道為1，大小為512*512的數組
        img = img.reshape(1, 1, img.shape[0], img.shape[1])
        # 轉為tensor
        img_tensor = torch.from_numpy(img)
        # 將tensor拷貝到device中，只用cpu就是拷貝到cpu中，用cuda就是拷貝到cuda中。
        img_tensor = img_tensor.to(device=device, dtype=torch.float32)
        # 預測
        pred = net(img_tensor)
        # 提取結果
        pred = np.array(pred.data.cpu()[0])[0]
        # 處理結果
        pred[pred >= 0.5] = 255
        pred[pred < 0.5] = 0
        # 保存圖片
        cv2.imwrite(save_res_path, pred)

運行完后，你可以在data/test目錄下，看到預測結果：

大功告成！

關于如何深度解析Pytorch中的UNet模型就分享到這里了，希望以上內容可以對大家有一定的幫助，可以學到更多知識。如果覺得文章不錯，可以把它分享出去讓更多的人看到。