Pytorch深度学习实战教程(三):UNet模型训练

2020年3月29日12:28:22 90 6,279 °C
摘要

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

Pytorch深度学习实战教程(三):UNet模型训练

一、前言

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

该系列文章的内容有:

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

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

二、项目背景

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

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

Pytorch深度学习实战教程(三):UNet模型训练

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

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

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

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

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

三、UNet训练

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

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

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

1、数据加载

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

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

数据已经备好,都在这里(Github):点击查看

如果 Github 下载速度慢,可以使用文末的百度链接下载数据集。

数据分为训练集和测试集,各30张,训练集有标签,测试集没有标签。

数据加载要做哪些处理,是根据任务和数据集而决定的,对于我们的分割任务,不用做太多处理,但由于数据量很少,仅30张,我们可以使用一些数据增强方法,来扩大我们的数据集。

Pytorch 给我们提供了一个方法,方便我们加载数据,我们可以使用这个框架,去加载我们的数据。看下伪代码:

这是一个标准的模板,我们就使用这个模板,来加载数据,定义标签,以及进行数据增强。

创建一个dataset.py文件,编写代码如下:

运行代码,你可以看到如下结果:

Pytorch深度学习实战教程(三):UNet模型训练

解释一下代码:

__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文件,编写如下代码:

创建unet_model.py文件,编写如下代码:

这样调整过后,网络的输出尺寸就与图片的输入尺寸相同了。

3、算法选择

选择什么 Loss 很重要,Loss 选择的好坏,都会影响算法拟合数据的效果。

选择什么 Loss 也是根据任务而决定的。我们今天的任务,只需要分割出细胞边缘,也就是一个很简单的二分类任务,所以我们可以使用 BCEWithLogitsLoss。

啥是 BCEWithLogitsLoss?BCEWithLogitsLoss 是 Pytorch 提供的用来计算二分类交叉熵的函数。

它的公式是:

Pytorch深度学习实战教程(三):UNet模型训练

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

具体的公式推导,可以看我的机器学习系列教程《机器学习实战教程(六):Logistic回归基础篇之梯度上升算法》,这里就不再累述。

目标函数,也就是 Loss 确定好了,怎么去优化这个目标呢?

最简单的方法就是,我们耳熟能详的梯度下降算法,逐渐逼近局部的极值。

但是这种简单的优化算法,求解速度慢,也就是想找到最优解,费劲儿。

各种优化算法,本质上其实都是梯度下降,例如最常规的 SGD,就是基于梯度下降改进的随机梯度下降算法,Momentum 就是引入了动量的 SGD,以指数衰减的形式累计历史梯度。

除了这些最基本的优化算法,还有自适应参数的优化算法。这类算法最大的特点就是,每个参数有不同的学习率,在整个学习过程中自动适应这些学习率,从而达到更好的收敛效果。

本文就是选择了一种自适应的优化算法 RMSProp。

由于篇幅有限,这里就不再扩展,讲解这个优化算法单写一篇都不够,要弄懂 RMSProp,你得先知道什么是 AdaGrad,因为 RMSProp 是基于 AdaGrad 的改进。

比 RMSProp 更高级的优化算法也有,比如大名鼎鼎的 Adam,它可以看做是修正后的Momentum+RMSProp 算法。

总之,对于初学者,你只要知道 RMSProp 是一种自适应的优化算法,比较高级就行了。

下面,我们就可以开始写训练UNet的代码了,创建 train.py 编写如下代码:

为了让工程更加清晰简洁,我们创建一个 model 文件夹,里面放模型相关的代码,也就是我们的网络结构代码,unet_parts.py 和 unet_model.py。

创建一个 utils 文件夹,里面放工具相关的代码,比如数据加载工具dataset.py。

这种模块化的管理,大大提高了代码的可维护性。

train.py 放在工程根目录即可,简单解释下代码。

由于数据就30张,我们就不分训练集和验证集了,我们保存训练集 loss 值最低的网络参数作为最佳模型参数。

如果都没有问题,你可以看到 loss 正在逐渐收敛。

Pytorch深度学习实战教程(三):UNet模型训练

四、预测

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

在工程根目录创建 predict.py 文件,编写如下代码:

运行完后,你可以在data/test目录下,看到预测结果:

Pytorch深度学习实战教程(三):UNet模型训练

大功告成!

五、最后

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

PS: 如果觉得本篇本章对您有所帮助,欢迎关注、评论、赞!

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

文件下载 ISBI数据集
下载地址
weinxin
微信公众号
分享技术,乐享生活:微信公众号搜索「JackCui-AI」关注一个在互联网摸爬滚打的潜行者。
Jack Cui

发表评论

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

目前评论:90   其中:访客  54   博主  36

    • avatar 学习 来自天朝的朋友 谷歌浏览器 Windows 10 重庆市 电信 4

      我记得最后一层应该有个Sigmoid,为啥没看到呢?