Pytorch 深度学习实战教程（三）：UNet模型训练，深度解析！

た入场券 2023-02-12 12:26 65阅读 0赞

![Pytorch深度学习实战教程（三）：UNet模型训练][Pytorch_UNet]

> 本文 GitHub [https://github.com/Jack-Cherish/PythonPark][https_github.com_Jack-Cherish_PythonPark] 已收录，有技术干货文章，整理的学习资料，一线大厂面试经验分享等，欢迎 Star 和 完善。

## 一、前言 ##

本文属于 Pytorch 深度学习语义分割系列教程。

该系列文章的内容有：

*  Pytorch 的基本使用
 *  语义分割算法讲解

PS：文中出现的所有代码，均可在我的 github 上下载，欢迎 Follow、Star：[点击查看][Link 1]

## 二、项目背景 ##

深度学习算法，无非就是我们解决一个问题的方法。选择什么样的网络去训练，进行什么样的预处理，采用什么Loss和优化方法，都是根据具体的任务而定的。

![aHR0cHM6Ly9jdWlqaWFodWEuY29tL3dwLWNvbnRlbnQvdXBsb2Fkcy8yMDIwLzAzL2RsLTE2LTEuZ2lm][]

所以，让我们先看一下今天的任务。

没错，就是 UNet 论文中的经典任务：医学图像分割。

选择它作为今天的任务，就是因为简单，好上手。

简单描述一个这个任务：如动图所示，给一张细胞结构图，我们要把每个细胞互相分割开来。

这个训练数据只有30张，分辨率为512x512，这些图片是果蝇的电镜图。

好了，任务介绍完毕，开始准备训练模型。

## 三、UNet训练 ##

想要训练一个深度学习模型，可以简单分为三个步骤：

*  数据加载：数据怎么加载，标签怎么定义，用什么数据增强方法，都是这一步进行。
 *  模型选择：模型我们已经准备好了，就是该系列上篇文章讲到的 UNet 网络。
 *  算法选择：算法选择也就是我们选什么 loss ，用什么优化算法。

每个步骤说的比较笼统，我们结合今天的医学图像分割任务，展开说明。

### 1、数据加载 ###

这一步，可以做很多事情，说白了，无非就是图片怎么加载，标签怎么定义，为了增加算法的鲁棒性或者增加数据集，可以做一些数据增强的操作。

既然是处理数据，那么我们先看下数据都是什么样的，再决定怎么处理。

数据已经备好，都在这里（Github）：[点击查看][Link 2]

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

运行代码，你可以看到如下结果：

![format_png][]

解释一下代码：

\_\_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语义分割网络][Pytorch_UNet 1]》讲解的 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 提供的用来计算二分类交叉熵的函数。

它的公式是：

![format_png 1][]

看过我机器学习系列教程的朋友，对这个公式一定不陌生，它就是 Logistic 回归的损失函数。它利用的是 Sigmoid 函数阈值在\[0,1\]这个特性来进行分类的。

具体的公式推导，可以看我的机器学习系列教程《[机器学习实战教程（六）：Logistic回归基础篇之梯度上升算法][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 正在逐渐收敛。

![format_png 2][]

## 四、预测 ##

模型训练好了，我们可以用它在测试集上看下效果。

在工程根目录创建 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模型训练][Pytorch_UNet 2]][Pytorch_UNet_Pytorch_UNet 2]

大功告成！

## 五、最后 ##

*  本文主要讲解了训练模型的三个步骤：数据加载、模型选择、算法选择。
 *  这是一个简单的例子，训练正常的视觉任务，要复杂很多。比如：在训练模型的时候，需要根据模型在验证集上的准确率选择保存哪个模型；需要支持tensorboard方便我们观察loss收敛情况等等。

> 点赞再看，养成习惯，微信公众号搜索**【JackCui-AI】**关注一个在互联网摸爬滚打的潜行者

![format_png 3][]

[Pytorch_UNet]: https://imgconvert.csdnimg.cn/aHR0cHM6Ly9jdWlqaWFodWEuY29tL3dwLWNvbnRlbnQvdXBsb2Fkcy8yMDIwLzAzL2RsLTE2LXRlc2UucG5n?x-oss-process=image/format,png
[https_github.com_Jack-Cherish_PythonPark]: https://github.com/Jack-Cherish/PythonPark
[Link 1]: https://github.com/Jack-Cherish/Deep-Learning/tree/master/Pytorch-Seg/lesson-2
[aHR0cHM6Ly9jdWlqaWFodWEuY29tL3dwLWNvbnRlbnQvdXBsb2Fkcy8yMDIwLzAzL2RsLTE2LTEuZ2lm]: https://imgconvert.csdnimg.cn/aHR0cHM6Ly9jdWlqaWFodWEuY29tL3dwLWNvbnRlbnQvdXBsb2Fkcy8yMDIwLzAzL2RsLTE2LTEuZ2lm
[Link 2]: https://github.com/Jack-Cherish/Deep-Learning/tree/master/Pytorch-Seg/lesson-2/data
[format_png]: https://imgconvert.csdnimg.cn/aHR0cHM6Ly9jdWlqaWFodWEuY29tL3dwLWNvbnRlbnQvdXBsb2Fkcy8yMDIwLzAzL2RsLTE2LTItMi5wbmc?x-oss-process=image/format,png
[Pytorch_UNet 1]: https://cuijiahua.com/blog/2019/12/dl-15.html
[format_png 1]: https://imgconvert.csdnimg.cn/aHR0cHM6Ly9jdWlqaWFodWEuY29tL3dwLWNvbnRlbnQvdXBsb2Fkcy8yMDIwLzAzL2RsLTE2LTMucG5n?x-oss-process=image/format,png
[Logistic]: https://cuijiahua.com/blog/2017/11/ml_6_logistic_1.html
[format_png 2]: https://imgconvert.csdnimg.cn/aHR0cHM6Ly9jdWlqaWFodWEuY29tL3dwLWNvbnRlbnQvdXBsb2Fkcy8yMDIwLzAzL2RsLTE2LTQucG5n?x-oss-process=image/format,png
[Pytorch_UNet 2]: https://imgconvert.csdnimg.cn/aHR0cHM6Ly9jdWlqaWFodWEuY29tL3dwLWNvbnRlbnQvdXBsb2Fkcy8yMDIwLzAzL2RsLTE2LTUucG5n?x-oss-process=image/format,png
[Pytorch_UNet_Pytorch_UNet 2]: https://cuijiahua.com/wp-content/uploads/2020/03/dl-16-5.png
[format_png 3]: https://imgconvert.csdnimg.cn/aHR0cHM6Ly9jdWlqaWFodWEuY29tL3dwLWNvbnRlbnQvdXBsb2Fkcy8yMDIwLzA1L2d6aC1sb2dvLXlhc3VvLmpwZw?x-oss-process=image/format,png