nnU-Net V2代码——生成splits_final.json文件

nnU-Net V2根据splits_final.json文件划分数据集，并支持自定义分割数据集。若未进行自定义分割数据集，则nnU-Net V2将自动生成默认的splits_final.json文件。

本文先阅读官方文档，查看如何自定义分割数据集，再阅读默认情况下生成splits_final.json文件的代码

nnU-Net V2官方文档说明

说明文档在nnUNet \ documentation \ manual_data_splits.md中

如何在 nnU-Net 中生成自定义数据分割

有时，nnU-Net 默认的 5 折交叉验证分割并不符合项目需求。也许您想运行 3 折交叉验证？或者您的训练案例不能随机分割，需要谨慎分层。别担心，nnU-Net 可以满足您的需求（它真的可以做到一切 ❤️）。

nnU-Net 使用的分割是在 nnUNetTrainer 的 do_split 函数中生成的。此函数会首先查找已存在的分割文件，如果不存在，则会创建一个。因此，如果您想影响分割，手动创建一个分割文件，使其被识别并使用，这是最佳选择！

do_split函数在nnUNet \ nnunetv2 \ training \ nnUNetTrainer \ nnUNetTrainer.py文件内

分割文件位于 nnUNet_preprocessed/DATASETXXX_NAME 文件夹中。因此，最佳实践是首先通过运行 nnUNetv2_plan_and_preproccess 来填充此文件夹。

运行nnUNetv2_plan_and_preproccess 命令后，nnUNet_preprocessed/DATASETXXX_NAME文件夹构建完毕

如果读者在nnUNet_raw/DATASETXXX_NAME文件夹内自定义了splits_final.json文件，则该文件在运行nnUNetv2_plan_and_preproccess 命令时，会被复制到nnUNet_preprocessed/DATASETXXX_NAME内

如果读者并未自定义该文件，它将在do_split函数（运行nnUNetv2_train命令时）中被创建

分割以 .json 文件形式存储。它们是一个简单的 Python 列表。该列表的长度即为包含的分割数量（在默认的 nnU-Net 中为 5）。每个列表条目是一个字典，键为 'train' 和 'val'。值又是包含每个集合中训练标识符的简单列表。为了说明这一点，我以 Dataset002 文件为例进行了一些操作：

In [1]: from batchgenerators.utilities.file_and_folder_operations import load_json

In [2]: splits = load_json('splits_final.json')

In [3]: len(splits)
Out[3]: 5

In [4]: splits[0].keys()
Out[4]: dict_keys(['train', 'val'])

In [5]: len(splits[0]['train'])
Out[5]: 16

In [6]: len(splits[0]['val'])
Out[6]: 4

In [7]: print(splits[0])
{'train': ['la_003', 'la_004', 'la_005', 'la_009', 'la_010', 'la_011', 'la_014', 'la_017', 'la_018', 'la_019', 'la_020', 'la_022', 'la_023', 'la_026', 'la_029', 'la_030'],
'val': ['la_007', 'la_016', 'la_021', 'la_024']}

如果您仍然不确定分割应该是什么样子，只需从 Medical Decathlon 下载一些参考数据集，开始训练（以生成分割），然后使用您喜欢的文本编辑器手动检查 .json 文件即可！

运行nnUNetv2_plan_and_preproccess命令和nnUNetv2_train命令，查看生成默认的splits_final.json文件

为了生成自定义分割，您只需要重现上述数据结构，并将其保存为 nnUNet_preprocessed/DATASETXXX_NAME 文件夹中的 splits_final.json 文件。然后像往常一样使用 nnUNetv2_train 等命令即可。

生成splits_final.json文件

具体实施生成的代码位于nnUNet \ nnunetv2 \ utilities \ crossval_split.py文件的generate_crossval_split 函数内

generate_crossval_split 函数参数：

参数名称	解释	默认值
train_identifiers	待划分的训练样本标识符列表，例如官方文档提供的代码里的['la_003', 'la_004', 'la_005']	List[str]类型
seed	随机种子	12345
n_splits	交叉验证的折数	5

过程：

# 初始化保存所有划分结果的列表
splits = []

利用sklearn库的KFold实施分割，实例化KFold

kfold = KFold(n_splits=n_splits, shuffle=True, random_state=seed)

创建for循环，利用kfold的split函数生成各个折包含的的元素索引

for i, (train_idx, test_idx) in enumerate(kfold.split(train_identifiers)):
    # 根据索引获取本折的训练集和验证集
    train_keys = np.array(train_identifiers)[train_idx]
    test_keys = np.array(train_identifiers)[test_idx]
    
    splits.append({})
    
    # 将本折的训练集和验证集存入字典
    splits[-1]['train'] = list(train_keys)
    splits[-1]['val'] = list(test_keys)

最后，返回划分结果

return splits