Pytorch 深度學習實戰教程(三):UNet模型訓練,深度解析!

Pytorch深度學習實戰教程(三):UNet模型訓練

本文 GitHub https://github.com/Jack-Cherish/PythonPark 已收錄,有技術乾貨文章,整理的學習資料,一線大廠面試經驗分享等,歡迎 Star 和 完善。

一、前言

本文屬於 Pytorch 深度學習語義分割系列教程。

該系列文章的內容有:

  • Pytorch 的基本使用
  • 語義分割算法講解

PS:文中出現的所有代碼,均可在我的 github 上下載,歡迎 Follow、Star:點擊查看

二、項目背景

深度學習算法,無非就是我們解決一個問題的方法。選擇什麼樣的網絡去訓練,進行什麼樣的預處理,採用什麼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模型訓練

大功告成!

五、最後

  • 本文主要講解了訓練模型的三個步驟:數據加載、模型選擇、算法選擇。
  • 這是一個簡單的例子,訓練正常的視覺任務,要複雜很多。比如:在訓練模型的時候,需要根據模型在驗證集上的準確率選擇保存哪個模型;需要支持tensorboard方便我們觀察loss收斂情況等等。

點贊再看,養成習慣,微信公衆號搜索【JackCui-AI】關注一個在互聯網摸爬滾打的潛行者

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章