[医学影像分割] nnunet 处理自己的数据集 + 并行训练 +训练时遇到的问题

version-1

import os
import shutil
import json

dataset_json = {
    "name" : "HNC Segmentation",
    "description" : "HNC pre-radiological treatment Segmentation",
    "tensorImageSize" : "3D",
    "modality": {
        "0": "T2"
    },
    "labels" : {
        "0": "background",
        "1": "GTVp",
        "2": "GTVn"
    },
    "file_ending": ".nii.gz",
    "numTraining" : 0,
    "numTest" : 0,
    "training" : [],
    "validation": [],
    "test" : []
}



results_path = "/mnt/home/pc/Chenq_team/HNC_dataset/nnunetv1_dataset/nnUNet_raw_data/Task018_HNC/"
dir_path = "/mnt/home/pc/Chenq_team/HNC_dataset/HNTSMRG24_train"

os.makedirs(results_path, exist_ok=True)
os.makedirs(os.path.join(results_path, "imagesTr"), exist_ok=True)
os.makedirs(os.path.join(results_path, "labelsTr"), exist_ok=True)
os.makedirs(os.path.join(results_path, "imagesTs"), exist_ok=True)


for dir in os.listdir(dir_path):
    pre_rt_dir_path = os.path.join(dir_path, dir, "preRT")
    pid = int(dir)
    img_new_path_json, mask_new_path = None, None
    for file in os.listdir(pre_rt_dir_path):
        if "T2" in file:
            img_old_path = os.path.join(pre_rt_dir_path, file)
            img_new_path = os.path.join(results_path, "imagesTr", f"preRT_{pid:03d}_0000.nii.gz")
            img_new_path_json = os.path.join(results_path, "imagesTr", f"preRT_{pid:03d}.nii.gz")
            shutil.copy(img_old_path, img_new_path)
        if "mask" in file:
            mask_old_path = os.path.join(pre_rt_dir_path, file)
            mask_new_path = os.path.join(results_path, "labelsTr", f"preRT_{pid:03d}.nii.gz")
            shutil.copy(mask_old_path, mask_new_path)
    assert img_new_path is not None and mask_new_path is not None, f"{pre_rt_dir_path} does not contain T2 or mask file"
    dataset_json["training"].append(
        {
            "image": img_new_path_json,
            "label": mask_new_path
        },
    )
    dataset_json["numTraining"] += 1

with open(os.path.join(results_path, "dataset.json"), "w") as f:
    json.dump(dataset_json, f, indent=4)

常见的命令

nnUNet_plan_and_preprocess -t 19 --verify_dataset_integrity  # 预处理数据集

自己划分数据集

更换split_final.pkl文件

version-1的并行训练

我没有找到相关的nnunet命令，但是nnunet-v1提供了run_training_DDP.py文件。训练命令如下:
其中MedNeXt没有提供DDP的trainer，自己抄一下nnunet的trainer就好了。记得替换sync_batchnorm，然后注意一下是半精度训练还是fp32。

export nnUNet_raw_data_base="/home/Guanjq/HNC_SEG_Data/nnunetv1_dataset/" # replace to your database
export nnUNet_preprocessed="/home/Guanjq/HNC_SEG_Data/nnunetv1_dataset/nnUNet_preprocessed/"
export RESULTS_FOLDER="/home/Guanjq/HNC_SEG_Data/nnunetv1_dataset/nnUNet_results/"

export TORCH_DISTRIBUTED_DEBUG="DETAIL"
export fold=3
export CUDA_VISIBLE_DEVICES=0,1 
export WORLD_SIZE=2
python3 -m torch.distributed.launch \
    --master_port 23331 \
    --nproc_per_node=2 \
    /home/Guanjq/Work/MedNeXt/nnunet_mednext/run/run_training_DDP.py \
    -network 3d_fullres \
    -network_trainer nnUNetTrainerV2_MedNeXt_S_kernel3_DDP \
    -task Task018_HNC_pre \
    -fold ${fold} \
    -p /home/Guanjq/HNC_SEG_Data/nnunetv1_dataset/nnUNet_preprocessed/Task018_HNC_pre/nnUNetPlansv2.1 \
    --fp32 

version-2

import os
import shutil
import json

dataset_json = {
    "name" : "HNC Segmentation",
    "description" : "HNC pre-radiological treatment Segmentation",
    "tensorImageSize" : "3D",
    "channel_names": {
        "0": "T2"
    },
    "labels" : {
        "background" : 0,
        "GTVp" : 1,
        "GTVn" : 2
    },
    "file_ending": ".nii.gz",
    "numTraining" : 0,
    "numTest" : 0,
    "training" : [],
    "validation": [],
    "test" : []
}



results_path = "/mnt/home/pc/Chenq_team/HNC_dataset/nnunet_dataset/raw/Dataset018_HNC/"
dir_path = "/mnt/home/pc/Chenq_team/HNC_dataset/HNTSMRG24_train"

os.makedirs(results_path, exist_ok=True)
os.makedirs(os.path.join(results_path, "imagesTr"), exist_ok=True)
os.makedirs(os.path.join(results_path, "labelsTr"), exist_ok=True)

for dir in os.listdir(dir_path):
    pre_rt_dir_path = os.path.join(dir_path, dir, "preRT")
    pid = int(dir)
    img_new_path, mask_new_path = None, None
    for file in os.listdir(pre_rt_dir_path):
        if "T2" in file:
            img_old_path = os.path.join(pre_rt_dir_path, file)
            img_new_path = os.path.join(results_path, "imagesTr", f"preRT_{pid:03d}_0000.nii.gz")
            shutil.copy(img_old_path, img_new_path)
        if "mask" in file:
            mask_old_path = os.path.join(pre_rt_dir_path, file)
            mask_new_path = os.path.join(results_path, "labelsTr", f"preRT_{pid:03d}.nii.gz")
            shutil.copy(mask_old_path, mask_new_path)
    assert img_new_path is not None and mask_new_path is not None, f"{pre_rt_dir_path} does not contain T2 or mask file"
    dataset_json["training"].append(
        {
            "image": img_new_path,
            "label": mask_new_path
        },
    )
    dataset_json["numTraining"] += 1

with open(os.path.join(results_path, "dataset.json"), "w") as f:
    json.dump(dataset_json, f, indent=4)

遇到的问题

1. RuntimeError: One or more background workers are no longer alive. Exiting. Please check the print statements above for the actual error message

在使用version-1训练时，报了这个错误，尝试了 export nnUNet_n_proc_DA="1" 都没办法解决问题。