2021-08-24

nnUNet代码学习——预处理代码学习

之前在实验室用了一段时间nnUNet代码，当时学习了一段时间，不过半途而废了。现在开始工作，感觉有必要将nnUNet的代码全部重新过一遍。那么就先从nnUNet预处理开始吧。首先先学习如何配置nnUNet。

图像之间的依赖关系：图像大小，spacing
nnUnet： 模块化的设计用来增改网络
如何在新数据集上运行nnUNet
nnUNet使用混合精度进行训练
注意pytorch>1.6、显卡、以及cudnn版本的协调
使用cpu可进行相关的数据增强
最好使用conda创建虚拟环境

nnUNet处理数据部分
nnUNet需要知道原始数据路径，需要设置三个环境变量：
1nnUNet_raw_data_base:储存原始数据以及crop之后的数据，该路径下最小有一个子文件夹
nnUNet_raw_data——每个任务一个对应的子文件夹；

同时原始数据需要遵守特定格式：
每个数据需要当作一个单独的任务Task，并且有Task名，同时使用3位整数进行区分。（自定义任务id>100）
    nnUNet_raw_data_base/nnUNet_raw_data/
        ├── Task001_BrainTumour
        ├── Task002_Heart
        ├── Task003_Liver
        ├── Task004_Hippocampus
        ├── Task005_Prostate
        ├── ...
 对于每个任务的子文件夹，也需要按照特定的格式进行存储；
        Task001_BrainTumour/
        ├── dataset.json （metedata of the dataset）
        ├── imagesTr ( 包含训练数据 )
        ├── (imagesTs) （可选, 包括测试数据）
        └── labelsTr （训练数据的GT）
 每个训练数据都有唯一的编号，同时训练数据以及label必须使用.nii.gz进行保存。
 GT中的数字必须连续，0代表背景，其余的代表各个类别。
 对于多模态数据，使用为四位数后缀进行区分。
 case命名规范为——case_identifier_XXXX.nii.gz. 其中XXXX代表模态后缀。
 多模态命名示例：
        nnUNet_raw_data_base/nnUNet_raw_data/Task001_BrainTumour/
        ├── dataset.json
        ├── imagesTr
        │?? ├── BRATS_001_0000.nii.gz
        │?? ├── BRATS_001_0001.nii.gz
        │?? ├── BRATS_001_0002.nii.gz
        │?? ├── BRATS_001_0003.nii.gz
        │?? ├── ...
        ├── imagesTs
        │?? ├── BRATS_485_0000.nii.gz
        │?? ├── BRATS_485_0001.nii.gz
        │?? ├── BRATS_485_0002.nii.gz
        │?? ├── BRATS_485_0003.nii.gz
        │?? ├── ...
        └── labelsTr
            ├── BRATS_001.nii.gz
            ├── ...
  单模态命名示例：
        nnUNet_raw_data_base/nnUNet_raw_data/Task002_Heart/
        ├── dataset.json
        ├── imagesTr
        │?? ├── la_003_0000.nii.gz
        │?? ├── la_004_0000.nii.gz
        │?? ├── ...
        ├── imagesTs
        │?? ├── la_001_0000.nii.gz
        │?? ├── la_002_0000.nii.gz
        │?? ├── ...
        └── labelsTr
            ├── la_003.nii.gz
            ├── la_004.nii.gz
            ├── ...
For each training case, all images must have the same geometry to ensure that their pixel arrays are aligned. Also make sure that all your data is co-registered!（这句话应该怎么理解）
 
 对于dataset.json文件，用来记录当前任务
{ 
 "name": "PROSTATE", 
 "description": "Prostate transitional zone and peripheral zone segmentation",
 "reference": "Radboud University, Nijmegen Medical Centre",
 "licence":"CC-BY-SA 4.0",
 "relase":"1.0 04/05/2018",
 "tensorImageSize": "4D",
 "modality": { 
   "0": "T2", 
   "1": "ADC"
 }, 
 "labels": { 
   "0": "background", 
   "1": "PZ", 
   "2": "TZ"
 }, 
 "numTraining": 32, 
 "numTest": 16,
 "training":[{"image":"./imagesTr/prostate_16.nii.gz","label":"./labelsTr/prostate_16.nii.gz"},{"image":"./imagesTr/prostate_04.nii.gz","label":"./labelsTr/prostate_04.nii.gz"},...], 
 "test": ["./imagesTs/prostate_08.nii.gz","./imagesTs/prostate_22.nii.gz","./imagesTs/prostate_30.nii.gz",...]
 }
 "training"和”testing“中需要列出所有的case路径，如果没有测试数据集。”test“：[]
 2. nnUNet_preprocessed 用来保存预处理的数据，训练时从该文件夹读取数据，最好是放在SSD上
 3. RESULTS_FOLDER： 保存预训练以及训练的模型

在Ubuntu上如何设置上面的三个路径：

    方法一：
    export nnUNet_preprocessed="/media/fabian/nnUNet_preprocessed"
    export RESULTS_FOLDER="/media/fabian/nnUNet_trained_models"
    然后输入 source /home/fabian/.bashrc 生效       
	方法二： 在终端中，输入       
    export nnUNet_raw_data_base="/media/fabian/nnUNet_raw"
    export nnUNet_preprocessed="/media/fabian/nnUNet_preprocessed"
    export RESULTS_FOLDER="/media/fabian/nnUNet_trained_models"
    不过这个只是临时起作用，终端关闭则需重设置。

实验计划以及预处理

nnUNet 根据数据特征（image sizes, voxel spacings, intensity information etc）等，创建3个UNET模型：

1、2D U-Net
2、3D U-Net that operated on full resolution images
3、3D U-Net cascade where the first U-Net creates a coarse segmentation map in downsampled images which is then refined by the second U-Net

通过输入，完成这一步：

nnUNet_plan_and_preprocess -t XXX --verify_dataset_integrity

XXX 是任务ID，可以是同输入多个任务ID；完成该命令行后，会在相关文件夹下，找到预处理过后的文件；同时会为模型生成plan文件（*.pkl），为后续训练做准备。只对训练文件有效

  --verify_dataset_integrity 在第一次跑的时候必须使用，用来检测数据是否能够使用nnUNET进行训练

模型训练

nnUNET使用5折交叉验证的方法。

对于小数据，级联UNET不会使用。

模型训练命令：

nnUNet_train CONFIGURATION TRAINER_CLASS_NAME TASK_NAME_OR_ID FOLD (additional options)
CONFIGURATION  UNET配置
TRAINER_CLASS_NAME  如何训练模型
TASK_NAME_OR_ID 任务名
FOLD 交叉验证
具体：
    2d模型：
    nnUNet_train 2d nnUNetTrainerV2 TaskXXX_MYTASK FOLD
    3d全尺寸
    nnUNet_train 3d_fullres nnUNetTrainerV2 TaskXXX_MYTASK FOLD
    3D级联
    nnUNet_train 3d_lowres nnUNetTrainerV2 TaskXXX_MYTASK FOLD
    nnUNet_train 3d_cascade_fullres nnUNetTrainerV2CascadeFullRes TaskXXX_MYTASK FOLD
    3D级联需要先完成低分辨率上的5折交叉验证之后，才能在全分辨率上进行训练。

训练好的模型将会保存在 RESULTS_FOLDER/nnUNet 文件夹下。保存文件夹名为：
nnUNet_preprocessed/CONFIGURATION/TaskXXX_MYTASKNAME/TRAINER_CLASS_NAME__PLANS_FILE_NAME/FOLD

举例：

    RESULTS_FOLDER/nnUNet/
    ├── 2d
    │?? └── Task02_Heart
    │??     └── nnUNetTrainerV2__nnUNetPlansv2.1
    │??         ├── fold_0
    │??         ├── fold_1
    │??         ├── fold_2
    │??         ├── fold_3
    │??         └── fold_4
    ├── 3d_cascade_fullres
    ├── 3d_fullres
    │?? └── Task02_Heart
    │??     └── nnUNetTrainerV2__nnUNetPlansv2.1
    │??         ├── fold_0
    │??         │?? ├── debug.json  （模型参数以及训练相关信息）
    │??         │?? ├── model_best.model  （得到最好的模型）
    │??         │?? ├── model_best.model.pkl
    │??         │?? ├── model_final_checkpoint.model  （最后一个epoch的模型）
    │??         │?? ├── model_final_checkpoint.model.pkl
    │??         │?? ├── network_architecture.pdf   （只有在安装hiddenlayer之后，才会有作用，模型结构）
    │??         │?? ├── progress.png      （训练、val以及评价指标（mDICE）的曲线）
    │??         │?? └── validation_raw     （训练结束后在验证集上的测试结果）
    │??         │??     ├── la_007.nii.gz
    │??         │??     ├── la_007.pkl
    │??         │??     ├── la_016.nii.gz
    │??         │??     ├── la_016.pkl
    │??         │??     ├── la_021.nii.gz
    │??         │??     ├── la_021.pkl
    │??         │??     ├── la_024.nii.gz
    │??         │??     ├── la_024.pkl
    │??         │??     ├── summary.json
    │??         │??     └── validation_args.json
    │??         ├── fold_1
    │??         ├── fold_2
    │??         ├── fold_3
    │??         └── fold_4
    └── 3d_lowres

多GPU训练(不推荐)

nnUNet支持两种多GPU实施方式，DataParallel（DP）以及Distributed Data Parallel (DDP，仅支持单主机)(两者差异)，DDP会更快但是如果使用pip install nnunet无法使用，需要单独编写额外的脚本。

因为损失函数、网络结构的变换，分布式训练不支持高分辨率UNet数据。

如果使用DP模式，使用到的命令为：

CUDA_VISIBLE_DEVICES=0,1,2... nnUNet_train_DP CONFIGURATION nnUNetTrainerV2_DP TASK_NAME_OR_ID FOLD -gpus GPUS --dbs

其中GPUs是用来训练的GPU数目

使用DDP模式进行训练，当前的目录必须是在nnUNet文件夹下面，可以使用进行训练

CUDA_VISIBLE_DEVICES=0,1,2... python -m torch.distributed.launch --master_port=XXXX --nproc_per_node=Y run/run_training_DDP.py CONFIGURATION nnUNetTrainerV2_DDP TASK_NAME_OR_ID FOLD --dbs

XXXX must be an open port for process-process communication (something like 4321 will do on most systems). Y is the number of GPUs you wish to use. Remember that we do not (yet) support distributed training across compute nodes. This all happens on the same system. Again, you can use CUDA_VISIBLE_DEVICES=0,1,2 to control what GPUs are used. If you run more than one DDP training on the same system (say you have 4 GPUs and you run two training with 2 GPUs each) you need to specify a different --master_port for each training!

使用多GPU训练，无法使用直接进行推理，在完成所有折的推理之后，在保存训练模型文件夹上运行
nnUNet_change_trainer_class之后，才能进行推理

找到最佳模型

使用以下命令进行验证

 nnUNet_train 2d nnUNetTrainerV2 Task102_mytask 0 -val --npz

在完成5折交叉验证之后，使用以下的命令会自动的找到最佳模型：

 nnUNet_find_best_configuration -m 2d 3d_fullres 3d_lowres 3d_cascade_fullres -t XXX --strict

可以手动调整需要搜索的配置

模型推理

测试数据所在文件夹必须满足如下配置：

所有的数据格式必须于训练数据相同，模态需要使用4位数字确定

多模态数据
     input_folder
     ├── prostate_03_0000.nii.gz
     ├── prostate_03_0001.nii.gz
     ├── prostate_05_0000.nii.gz
     ├── prostate_05_0001.nii.gz
     ├── prostate_08_0000.nii.gz
     ├── prostate_08_0001.nii.gz
     ├── ...
单模态数据
     imagesTs
     ├── la_001_0000.nii.gz
     ├── la_002_0000.nii.gz
     ├── la_006_0000.nii.gz
     ├── ...

如果需要手动对特定配置进行推理，使用：

nnUNet_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -t TASK_NAME_OR_ID -m CONFIGURATION --save_npz

save_npz 会保存经过softmax后的预测结果，会占用大量内存

同时需要对每个配置进行指定保存文件夹

如果需要进行集成

nnUNet_ensemble -f FOLDER1 FOLDER2 ... -o OUTPUT_FOLDER -pp POSTPROCESSING_FILE

进行集成，还需要指定一个文件说明如何进行后处理，同时需要保存在指定的文件路径(RESULTS_FOLDER/nnUNet/CONFIGURATION/TaskXXX_MYTASK/TRAINER_CLASS_NAME__PLANS_FILE_IDENTIFIER/postprocessing.json or RESULTS_FOLDER/nnUNet/ensembles/TaskXXX_MYTASK/ensemble_X__Y__Z–X__Y__Z/postprocessing.json). 也可以通过忽略PP选项不提供文件，不进行后处理

推荐对5折得到的模型都进行inference

使用预训练模型进行推理

nnUNet预训练模型都是公开的，下载预训练模型会覆盖相关的配置、trainer以及plans.

使用以下命令可以得到相关的公开模型

nnUNet_print_available_pretrained_models

使用以下命令可以下载相关模型

nnUNet_download_pretrained_model Task029_LiTS

完成模型下载之后，可以直接进行推理