unet模型训练

本文由简悦 SimpRead 转码，原文地址 cuijiahua.com

本文进行实战学习，针对医学图像分割任务，讲解了训练模型的三个步骤：数据加载、模型选择、算法选择。.

摘要

本文进行实战学习，针对医学图像分割任务，讲解了训练模型的三个步骤：数据加载、模型选择、算法选择。

一、前言

本文属于 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 目录下，看到预测结果：

大功告成！

五、最后

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

PS：如果觉得本篇本章对您有所帮助，欢迎关注、评论、赞！

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