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

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

一、項(xiàng)目背景

深度學(xué)習(xí)算法，無(wú)非就是我們解決一個(gè)問(wèn)題的方法。選擇什么樣的網(wǎng)絡(luò)去訓(xùn)練，進(jìn)行什么樣的預(yù)處理，采用什么Loss和優(yōu)化方法，都是根據(jù)具體的任務(wù)而定的。

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

沒(méi)錯(cuò)，就是 UNet 論文中的經(jīng)典任務(wù)：醫(yī)學(xué)圖像分割。

選擇它作為今天的任務(wù)，就是因?yàn)楹?jiǎn)單，好上手。

簡(jiǎn)單描述一個(gè)這個(gè)任務(wù)：如動(dòng)圖所示，給一張細(xì)胞結(jié)構(gòu)圖，我們要把每個(gè)細(xì)胞互相分割開(kāi)來(lái)。

這個(gè)訓(xùn)練數(shù)據(jù)只有30張，分辨率為512x512，這些圖片是果蠅的電鏡圖。

好了，任務(wù)介紹完畢，開(kāi)始準(zhǔn)備訓(xùn)練模型。

二、UNet訓(xùn)練

想要訓(xùn)練一個(gè)深度學(xué)習(xí)模型，可以簡(jiǎn)單分為三個(gè)步驟：

• 數(shù)據(jù)加載：數(shù)據(jù)怎么加載，標(biāo)簽怎么定義，用什么數(shù)據(jù)增強(qiáng)方法，都是這一步進(jìn)行。

• 模型選擇：模型我們已經(jīng)準(zhǔn)備好了，就是該系列上篇文章講到的 UNet 網(wǎng)絡(luò)。

• 算法選擇：算法選擇也就是我們選什么 loss ，用什么優(yōu)化算法。

每個(gè)步驟說(shuō)的比較籠統(tǒng)，我們結(jié)合今天的醫(yī)學(xué)圖像分割任務(wù)，展開(kāi)說(shuō)明。

1、數(shù)據(jù)加載

這一步，可以做很多事情，說(shuō)白了，無(wú)非就是圖片怎么加載，標(biāo)簽怎么定義，為了增加算法的魯棒性或者增加數(shù)據(jù)集，可以做一些數(shù)據(jù)增強(qiáng)的操作。

既然是處理數(shù)據(jù)，那么我們先看下數(shù)據(jù)都是什么樣的，再?zèng)Q定怎么處理。

數(shù)據(jù)已經(jīng)備好，都在這里（Github）：點(diǎn)擊查看

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

數(shù)據(jù)分為訓(xùn)練集和測(cè)試集，各30張，訓(xùn)練集有標(biāo)簽，測(cè)試集沒(méi)有標(biāo)簽。

數(shù)據(jù)加載要做哪些處理，是根據(jù)任務(wù)和數(shù)據(jù)集而決定的，對(duì)于我們的分割任務(wù)，不用做太多處理，但由于數(shù)據(jù)量很少，僅30張，我們可以使用一些數(shù)據(jù)增強(qiáng)方法，來(lái)擴(kuò)大我們的數(shù)據(jù)集。

Pytorch 給我們提供了一個(gè)方法，方便我們加載數(shù)據(jù)，我們可以使用這個(gè)框架，去加載我們的數(shù)據(jù)?？聪聜未a：

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

這是一個(gè)標(biāo)準(zhǔn)的模板，我們就使用這個(gè)模板，來(lái)加載數(shù)據(jù)，定義標(biāo)簽，以及進(jìn)行數(shù)據(jù)增強(qiáng)。

創(chuàng)建一個(gè)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):
        # 初始化函數(shù)，讀取所有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進(jìn)行數(shù)據(jù)增強(qiáng)，filpCode為1水平翻轉(zhuǎn)，0垂直翻轉(zhuǎn)，-1水平+垂直翻轉(zhuǎn)
        flip = cv2.flip(image, flipCode)
        return flip
        
    def __getitem__(self, index):
        # 根據(jù)index讀取圖片
        image_path = self.imgs_path[index]
        # 根據(jù)image_path生成label_path
        label_path = image_path.replace('image', 'label')
        # 讀取訓(xùn)練圖片和標(biāo)簽圖片
        image = cv2.imread(image_path)
        label = cv2.imread(label_path)
        # 將數(shù)據(jù)轉(zhuǎn)為單通道的圖片
        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])
        # 處理標(biāo)簽，將像素值為255的改為1
        if label.max() > 1:
            label = label / 255
        # 隨機(jī)進(jìn)行數(shù)據(jù)增強(qiáng)，為2時(shí)不做處理
        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):
        # 返回訓(xùn)練集大小
        return len(self.imgs_path)

    
if __name__ == "__main__":
    isbi_dataset = ISBI_Loader("data/train/")
    print("數(shù)據(jù)個(gè)數(shù)：", 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)

運(yùn)行代碼，你可以看到如下結(jié)果：

解釋一下代碼：

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

__len__函數(shù)可以返回?cái)?shù)據(jù)的多少，這個(gè)類實(shí)例化后，通過(guò)len()函數(shù)調(diào)用。

__getitem__函數(shù)是數(shù)據(jù)獲取函數(shù)，在這個(gè)函數(shù)里你可以寫數(shù)據(jù)怎么讀，怎么處理，并且可以一些數(shù)據(jù)預(yù)處理、數(shù)據(jù)增強(qiáng)都可以在這里進(jìn)行。我這里的處理很簡(jiǎn)單，只是將圖片讀取，并處理成單通道圖片。同時(shí)，因?yàn)?label 的圖片像素點(diǎn)是0和255，因此需要除以255，變成0和1。同時(shí)，隨機(jī)進(jìn)行了數(shù)據(jù)增強(qiáng)。

augment函數(shù)是定義的數(shù)據(jù)增強(qiáng)函數(shù)，怎么處理都行，我這里只是進(jìn)行了簡(jiǎn)單的旋轉(zhuǎn)操作。

在這個(gè)類中，你不用進(jìn)行一些打亂數(shù)據(jù)集的操作，也不用管怎么按照 batchsize 讀取數(shù)據(jù)。因?yàn)閷?shí)例化這個(gè)類后，我們可以用 torch.utils.data.DataLoader 方法指定 batchsize 的大小，決定是否打亂數(shù)據(jù)。

Pytorch 提供給給我們的 DataLoader 很強(qiáng)大，我們甚至可以指定使用多少個(gè)進(jìn)程加載數(shù)據(jù)，數(shù)據(jù)是否加載到 CUDA 內(nèi)存中等高級(jí)用法，本文不涉及，就不再展開(kāi)講解了。

2、模型選擇

模型我們已經(jīng)選擇完了，就用上篇文章《Pytorch深度學(xué)習(xí)實(shí)戰(zhàn)教程（二）：UNet語(yǔ)義分割網(wǎng)絡(luò)》講解的 UNet 網(wǎng)絡(luò)結(jié)構(gòu)。

但是我們需要對(duì)網(wǎng)絡(luò)進(jìn)行微調(diào)，完全按照論文的結(jié)構(gòu)，模型輸出的尺寸會(huì)稍微小于圖片輸入的尺寸，如果使用論文的網(wǎng)絡(luò)結(jié)構(gòu)需要在結(jié)果輸出后，做一個(gè) resize 操作。為了省去這一步，我們可以修改網(wǎng)絡(luò)，使網(wǎng)絡(luò)的輸出尺寸正好等于圖片的輸入尺寸。

創(chuàng)建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)

創(chuàng)建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)

這樣調(diào)整過(guò)后，網(wǎng)絡(luò)的輸出尺寸就與圖片的輸入尺寸相同了。

3、算法選擇

選擇什么 Loss 很重要，Loss 選擇的好壞，都會(huì)影響算法擬合數(shù)據(jù)的效果。

選擇什么 Loss 也是根據(jù)任務(wù)而決定的。我們今天的任務(wù)，只需要分割出細(xì)胞邊緣，也就是一個(gè)很簡(jiǎn)單的二分類任務(wù)，所以我們可以使用 BCEWithLogitsLoss。

啥是 BCEWithLogitsLoss？BCEWithLogitsLoss 是 Pytorch 提供的用來(lái)計(jì)算二分類交叉熵的函數(shù)。

它的公式是：

看過(guò)我機(jī)器學(xué)習(xí)系列教程的朋友，對(duì)這個(gè)公式一定不陌生，它就是 Logistic 回歸的損失函數(shù)。它利用的是 Sigmoid 函數(shù)閾值在[0,1]這個(gè)特性來(lái)進(jìn)行分類的。

具體的公式推導(dǎo)，可以看我的機(jī)器學(xué)習(xí)系列教程《機(jī)器學(xué)習(xí)實(shí)戰(zhàn)教程（六）：Logistic回歸基礎(chǔ)篇之梯度上升算法》，這里就不再累述。

目標(biāo)函數(shù)，也就是 Loss 確定好了，怎么去優(yōu)化這個(gè)目標(biāo)呢？

最簡(jiǎn)單的方法就是，我們耳熟能詳?shù)奶荻认陆邓惴?，逐漸逼近局部的極值。

但是這種簡(jiǎn)單的優(yōu)化算法，求解速度慢，也就是想找到最優(yōu)解，費(fèi)勁兒。

各種優(yōu)化算法，本質(zhì)上其實(shí)都是梯度下降，例如最常規(guī)的 SGD，就是基于梯度下降改進(jìn)的隨機(jī)梯度下降算法，Momentum 就是引入了動(dòng)量的 SGD，以指數(shù)衰減的形式累計(jì)歷史梯度。

除了這些最基本的優(yōu)化算法，還有自適應(yīng)參數(shù)的優(yōu)化算法。這類算法最大的特點(diǎn)就是，每個(gè)參數(shù)有不同的學(xué)習(xí)率，在整個(gè)學(xué)習(xí)過(guò)程中自動(dòng)適應(yīng)這些學(xué)習(xí)率，從而達(dá)到更好的收斂效果。

本文就是選擇了一種自適應(yīng)的優(yōu)化算法 RMSProp。

由于篇幅有限，這里就不再擴(kuò)展，講解這個(gè)優(yōu)化算法單寫一篇都不夠，要弄懂 RMSProp，你得先知道什么是 AdaGrad，因?yàn)?RMSProp 是基于 AdaGrad 的改進(jìn)。

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

總之，對(duì)于初學(xué)者，你只要知道 RMSProp 是一種自適應(yīng)的優(yōu)化算法，比較高級(jí)就行了。

下面，我們就可以開(kāi)始寫訓(xùn)練UNet的代碼了，創(chuàng)建 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):
    # 加載訓(xùn)練集
    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統(tǒng)計(jì)，初始化為正無(wú)窮
    best_loss = float('inf')
    # 訓(xùn)練epochs次
    for epoch in range(epochs):
        # 訓(xùn)練模式
        net.train()
        # 按照batch_size開(kāi)始訓(xùn)練
        for image, label in train_loader:
            optimizer.zero_grad()
            # 將數(shù)據(jù)拷貝到device中
            image = image.to(device=device, dtype=torch.float32)
            label = label.to(device=device, dtype=torch.float32)
            # 使用網(wǎng)絡(luò)參數(shù)，輸出預(yù)測(cè)結(jié)果
            pred = net(image)
            # 計(jì)算loss
            loss = criterion(pred, label)
            print('Loss/train', loss.item())
            # 保存loss值最小的網(wǎng)絡(luò)參數(shù)
            if loss < best_loss:
                best_loss = loss
                torch.save(net.state_dict(), 'best_model.pth')
            # 更新參數(shù)
            loss.backward()
            optimizer.step()

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

為了讓工程更加清晰簡(jiǎn)潔，我們創(chuàng)建一個(gè) model 文件夾，里面放模型相關(guān)的代碼，也就是我們的網(wǎng)絡(luò)結(jié)構(gòu)代碼，unet_parts.py 和 unet_model.py。

創(chuàng)建一個(gè) utils 文件夾，里面放工具相關(guān)的代碼，比如數(shù)據(jù)加載工具dataset.py。

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

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

由于數(shù)據(jù)就30張，我們就不分訓(xùn)練集和驗(yàn)證集了，我們保存訓(xùn)練集 loss 值最低的網(wǎng)絡(luò)參數(shù)作為最佳模型參數(shù)。

如果都沒(méi)有問(wèn)題，你可以看到 loss 正在逐漸收斂。

三、預(yù)測(cè)

模型訓(xùn)練好了，我們可以用它在測(cè)試集上看下效果。

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

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

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

運(yùn)行完后，你可以在data/test目錄下，看到預(yù)測(cè)結(jié)果：

大功告成！

關(guān)于如何深度解析Pytorch中的UNet模型就分享到這里了，希望以上內(nèi)容可以對(duì)大家有一定的幫助，可以學(xué)到更多知識(shí)。如果覺(jué)得文章不錯(cuò)，可以把它分享出去讓更多的人看到。