1
前言
本文属于 Pytorch 深度学习语义分割系列教程。
该系列文章的内容有:
Pytorch 的基本使用
语义分割算法讲解
由于微信不允许外部链接,你需要点击页尾左下角的“阅读原文”,才能访问文中的链接,文中的所有外部链接都已使用蓝色字体标记。
2
项目背景
深度学习算法,无非就是我们解决一个问题的方法。选择什么样的网络去训练,进行什么样的预处理,采用什么Loss和优化方法,都是根据具体的任务而定的。
所以,让我们先看一下今天的任务。
没错,就是 UNet 论文中的经典任务:医学图像分割。
选择它作为今天的任务,就是因为简单,好上手。
简单描述一个这个任务:如动图所示,给一张细胞结构图,我们要把每个细胞互相分割开来。
这个训练数据只有30张,分辨率为512x512,这些图片是果蝇的电镜图。
好了,任务介绍完毕,开始准备训练模型。
3
UNet训练
想要训练一个深度学习模型,可以简单分为三个步骤:
数据加载:数据怎么加载,标签怎么定义,用什么数据增强方法,都是这一步进行。
模型选择:模型我们已经准备好了,就是该系列上篇文章讲到的 UNet 网络。
算法选择:算法选择也就是我们选什么 loss ,用什么优化算法。
每个步骤说的比较笼统,我们结合今天的医学图像分割任务,展开说明。
1、数据加载
这一步,可以做很多事情,说白了,无非就是图片怎么加载,标签怎么定义,为了增加算法的鲁棒性或者增加数据集,可以做一些数据增强的操作。
既然是处理数据,那么我们先看下数据都是什么样的,再决定怎么处理。
数据已经备好,都在这里:(点原文连接了解更多)
https://github.com/Jack-Cherish/Deep-Learning/tree/master/Pytorch-Seg/lesson-2/data
数据分为训练集和测试集,各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、模型选择
模型我们已经选择为 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]这个特性来进行分类的。
目标函数,也就是 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 正在逐渐收敛。
4
预测
模型训练好了,我们可以用它在测试集上看下效果。
在工程根目录创建 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目录下,看到预测结果:
大功告成!
5
最后
本文主要讲解了训练模型的三个步骤:数据加载、模型选择、算法选择。
这是一个简单的例子,训练正常的视觉任务,要复杂很多。比如:在训练模型的时候,需要根据模型在验证集上的准确率选择保存哪个模型;需要支持tensorboard方便我们观察loss收敛情况等等。
PS:如果觉得本篇本章对您有所帮助,欢迎关注、评论、赞!
更多精彩内容(请点击图片进行阅读)
公众号:AI蜗牛车
保持谦逊、保持自律、保持进步
个人微信
备注:昵称+学校/公司+方向
如果没有备注不拉群!
拉你进AI蜗牛车交流群
点个在看,么么哒!