3D-UNet的Pytorch实现

3D-UNet的Pytorch实现本文主要介绍3DUNet网络，及其在LiTS2017肝脏肿瘤数据集上训练的Pytorch实现代码。GitHub地址：添加链接描述LiTS2017数据集链接：添加链接描述提取码：hfl8 （++||…==’’。。。_）一.3DUNet简介最近重新整理了一下关于3DUNet网络原理及代码，这个网络其实和2DUNet区别不大，简单说可以理解为2d卷积换为了3d卷积。整

文章共1,807字 · 阅读需要大约7分钟

一键AI生成摘要，助你高效阅读

问答

pawerd

14101人浏览 · 2020-12-08 10:05:16

pawerd · 2020-12-08 10:05:16 发布

3D-UNet的Pytorch实现

本文主要介绍3DUNet网络，及其在LiTS2017肝脏肿瘤数据集上训练的Pytorch实现代码。

GitHub地址：

(https://github.com/lee-zq/3DUNet-Pytorch%20%E2%80%8B%20github.com)

一.3DUNet简介
最近重新整理了一下关于3DUNet网络原理及代码，这个网络其实和2DUNet区别不大，简单说可以理解为2d卷积换为了3d卷积。整体上没有什么创新，但可以基于一套完整的3DUNet代码（包括预处理、训练、可视化、测试等等）可以简化很多工作，在此基础上实现更多的细节改进，比如设计替换最新模块等等。对比下图中的2dUNet和3dUNet网络结构：

3DUNet网络结构
在这里插入图片描述

2DUNet网络结构
在这里插入图片描述

除了一些超参数设置不同，以及2d和3d卷积的区别，两者设计思路几乎完全一样。所以在网络结构上没啥要说的。

二. 3DUNet的Pytorch实现

本文的3DUNet代码主要参考了这个项目（here），修改了一些bug并进行了代码重构和梳理。可以直接访问下面的github仓库链接download并按照readme步骤使用：

https://github.com/lee-zq/3DUNet-Pytorch

github.com
在这里我也再梳理一下代码结构和设计思路，以及使用方法。

准备工作
首先下载代码：

git clone https://github.com/lee-zq/3DUNet-Pytorch.git
下载的代码结构和对应的功能如下：

│  .gitignore
│  config.py              # 超参数配置
│  README.md              # 使用方法介绍
│  train.py               # 模型训练与验证函数  （主函数）
│  train_faster.py        # 针对加载数据太慢，更换了dataloader的模型训练函数   （主函数）
│  test.py                # 针对每个测试样本分patch进行推理并拼接为分割结果
│
├─dataset
│  │  dataset_lits.py    # 对LiTS数据集设计的dataset类，train.py调用
│  │  dataset_lits_faster.py   # 对LiTS数据集设计的dataset类，train_faster.py调用
├─models
│  │  Unet.py            # 2D、3DUNet网络模型
│  │
│  ├─nn                 # 网络公用模块存放位置，目前还未更新
│
├─output                 # trained models 都保存在这里面
│      readme.md          
│
├─preprocess             # 预处理模块，主要对原始数据集进行初步处理，和项目内其他代码完全隔离。如格式转换、切片提取等等             
│      preprocess_LiTS.py # 针对LiTS2017数据集设计的预处理方案
│
├─utils                  # 其他模块
│  │  common.py          # 常用函数包
│  │  init_util.py       # 初始化相关包
│  │  logger.py          # 日志生成函数
│  │  metrics.py         # 评价指标

然后从文章开头的百度云链接下载LiTS2017数据集。

此外还有python环境配置要求：

pytorch >= 1.1.0
torchvision
SimpleITK
Tensorboard
Scipy

由上面的项目结构基本就可以知道代码运行的步骤了，如下：

预处理步骤

将下载的LiTS数据集解压至任一目录（例如./raw_dataset/），这里建议将batch1，将batch2（28~130）作为训练集。然后分别将训练集和测试集里面的volum数据和segmentation数据分别放入该目录下的data文件夹和label文件夹（若计算资源有限或测试代码，可先提取一部分数据）。如下：

raw_dataset:
    ├── LiTS_batch1  # (0~27)
    │   ├── data
    │   │   ├── volume-0.nii
    │   │   ├── volume-10.nii ...
    │   └── label
    │       ├── segmentation-0.nii
    │       ├── segmentation-10.nii ...
    │       
    ├── LiTS_batch2 # (28~130)
    │   ├── data
    │   │   ├── volume-28.nii
    │   │   ├── volume-29.nii ...
    │   └── label
    │       ├── segmentation-28.nii
    │       ├── segmentation-29.nii ...

然后在 preprocess/preprocess_LiTS.py文件中更改预处理输入、输出目录：

raw_dataset_path = './raw_dataset/LiTS_batch2/'  # 输入数据集路径
fixed_dataset_path = './fixed_data/'  # 预处理后的数据集的输出路径

执行 python preprocess/preprocess_LiTS.py ，运行完后在 fixed_dataset_path=’./fixed_data/’ 目录下可生成如下内容：

│  train_name_list.txt    # 训练集文件名 list
│  val_name_list.txt      # 验证集文件名 list
│
├─data               # 预处理后的data数据
│      volume-28.nii
│      volume-29.nii
│      volume-30.nii
│      ...
└─label              # 预处理后的标签数据
        segmentation-28.nii
        segmentation-29.nii
        segmentation-30.nii
        ...

到这里，数据预处理完成。预处理的内容可通过源代码了解，我也加了中文注释。

模型训练

首先在 config.py中修改超参数–dataset_path为我们预处理后的数据根目录：./fixed_data ，其他参数也可以根据注释自行修改。当然你也可以在下面运行训练命令的时候指定参数去覆盖默认参数。

然后运行 python train.py --save model_name

这样训练即开始，到指定迭代次数，就会在output/model_name下生成保存的模型文件和日志文件，你可以通过 tensorboard --logdir ./output/model_name在浏览器查看训练过程中dice和loss。

训练模型时，大概率会出现GPU利用率不满/跳变的情况，这时建议使用 train_faster.py替换 train.py进行训练，可以实现从一个输入样本提取一个batch进行训练（通过调用dataset_lits_faster.py 实现），这样速度会快很多，但可能导致收敛较慢。

PS：一些问题
采用的数据集包含三类标签：背景、肝脏、肝肿瘤。肿瘤附着在肝脏上，体积很小。所以，直接按照三类来进行训练，会导致肿瘤分割效果较差，这个可以通过实验结果验证。所以对此类问题我们一般都是先分割或检测肝脏得到肝脏ROI，然后在此ROI内完成肿瘤分割。
保存模型的机制还需要改进，后续会增加保存验证指标最高和最新epoch得到的模型。

参考：

(https://github.com/panxiaobai/lits_pytorch)
(https://zhuanlan.zhihu.com/p/104854615)