nnUNet V2代码——生成nnUNetPlans.json（一）

前文请见nnUNetv2_plan_and_preprocess命令

阅读nnUNet\nnunetv2\experiment_planning\experiment_planners\default_experiment_planner.py

文件包含一个ExperimentPlanner类和一个_maybe_copy_splits_file函数。

在ExperimentPlanner函数内涉及的其他函数都在文章后半部分说明。

生成nnUNetPlans.json共三篇：

nnUNet V2代码——生成nnUNetPlans.json（一）

nnUNet V2代码——生成nnUNetPlans.json（二）

nnUNet V2代码——生成nnUNetPlans.json（三）

阅读nnUNet\nnunetv2\experiment_planning\experiment_planners\default_experiment_planner.py

文件包含一个ExperimentPlanner类和一个_maybe_copy_splits_file函数。

在ExperimentPlanner函数内涉及的其他函数都在文章后半部分说明。

一. __init__函数

1. 参数

• dataset_name_or_id: 数据集名称或ID
• gpu_memory_target_in_gb：预期GPU内存（单位：GB），在运行nnUNetv2_plan_and_preprocess命令时设置
• preprocessor_name：预处理器的名称，默认值为 DefaultPreprocessor。
• plans_name：训练计划的名称，默认值为 nnUNetPlans。
• overwrite_target_spacing：重写目标体素间距，只影响 3d_fullres 和 3d_cascade两种情况，默认是None。
• suppress_transpose：是否强制不转置，默认值为 False。

2. 过程

依次获取数据集名称、定义是否强制不转置、获取nnUNet_raw下的数据集路径、获取nnUNet_preprocessed下的数据集路径、读取dataset.json文件。

self.dataset_name = maybe_convert_to_dataset_name(dataset_name_or_id)
self.suppress_transpose = suppress_transpose
self.raw_dataset_folder = join(nnUNet_raw, self.dataset_name)
preprocessed_folder = join(nnUNet_preprocessed, self.dataset_name)
self.dataset_json = load_json(join(self.raw_dataset_folder, 'dataset.json'))
self.dataset = get_filenames_of_train_images_and_targets(self.raw_dataset_folder, self.dataset_json)

载入dataset_fingerprint.json文件，如果其不存在则报错。

设置各向异性阈值以及UNet结构的相关信息

self.anisotropy_threshold = ANISO_THRESHOLD
self.UNet_base_num_features = 32
self.UNet_class = PlainConvUNet

上面的 ANISO_THRESHOLD 值为 3
它的作用是确定何时将样本视为各向异性
当一个轴的体素间距大于其他轴体素间距的ANISO_THRESHOLD倍时
判定该数据集具有各向异性
各向异性会影响训练效果
nnUNet作者会在后续函数中优化调整具有各向异性的轴

注释机翻：下面代码的两个数字实际上是任意的，设定它们是为了尽可能地重现 nnU-Net v1 的配置

self.UNet_reference_val_3d = 560000000  # 455600128  550000000
self.UNet_reference_val_2d = 85000000  # 83252480

设置UNet结构和训练的相关信息

self.UNet_reference_com_nfeatures = 32
self.UNet_reference_val_corresp_GB = 8
self.UNet_reference_val_corresp_bs_2d = 12
self.UNet_reference_val_corresp_bs_3d = 2
self.UNet_featuremap_min_edge_length = 4
self.UNet_blocks_per_stage_encoder = (2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2)
self.UNet_blocks_per_stage_decoder = (2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2)
self.UNet_min_batch_size = 2
self.UNet_max_features_2d = 512
self.UNet_max_features_3d = 320
self.max_dataset_covered = 0.05  # 我们限制批次大小，以便在单次前向/反向传播中最多只能看到 5% 的数据集

接下来设置预期显存大小（单位是GB）以及一个阈值
注释机翻：如果在3d_fullres下的patch块大小小于image（中位数）大小的25%（此处设置的阈值），那么我们也需要一个3d_lowres的配置（这个阈值在plan_experiment函数中使用，控制3d_lowres的体素间距）。

self.UNet_vram_target_GB = gpu_memory_target_in_gb
self.lowres_creation_threshold = 0.25

继续设置预处理器名称，训练计划名称，用户自定义的体素间距（self.overwrite_target_spacing变量），两个断言是确保用户知晓自己的设置并合理的设置体素间距

self.preprocessor_name = preprocessor_name
self.plans_identifier = plans_name
self.overwrite_target_spacing = overwrite_target_spacing
assert overwrite_target_spacing is None or len(overwrite_target_spacing), 'if overwrite_target_spacing is ' \
                                                                            'used then three floats must be ' \
                                                                            'given (as list or tuple)'
assert overwrite_target_spacing is None or all([isinstance(i, float) for i in overwrite_target_spacing]), \
    'if overwrite_target_spacing is used then three floats must be given (as list or tuple)'

最后判断用户是否有自定义分割数据集的配置文件（见documentation/manual_data_splits.md文件），如果存在（用户自己设置的分割数据集配置文件），就将其复制到nnUNet_preprocessed文件夹下。

self.plans = None
if isfile(join(self.raw_dataset_folder, 'splits_final.json')):
    _maybe_copy_splits_file(join(self.raw_dataset_folder, 'splits_final.json'),
                            join(preprocessed_folder, 'splits_final.json'))

二. determine_fullres_target_spacing函数

1. 参数

无

2. 过程

本函数从名字可以看出，是确定体素间距的
首先判断用户是否自定义了体素间距（神经网络输入图像的体素间距），如果设置了就返回用户自定义的体素间距，否则执行后续代码程序

if self.overwrite_target_spacing is not None:
    return np.array(self.overwrite_target_spacing)

overwrite_target_spacing变量是在执行nnUNetv2_plan_and_preprocess时输入的，默认是None（nnUNet作者特别说明：除非用户知道自己设置此变量是在做什么，否则默认即可；如果用户设置了目标体素间距，最好同时设置overwrite_target_spacing用于生成不同名称的Plans.json文件，和默认生成的Plans.json区分开）

接下来是获取dataset_fingerprint.json文件内的”spacings“和”shapes_after_crop“两组数据。dataset_fingerprint.json的生成见这里。

spacings = np.vstack(self.dataset_fingerprint['spacings'])
sizes = self.dataset_fingerprint['shapes_after_crop']

查找spacings和sizes每一列（各个轴）的中位数，target存储每个轴体素间距的中位数，target_size存储每个轴图像大小的中位数：

target = np.percentile(spacings, 50, 0)
target_size = np.percentile(np.vstack(sizes), 50, 0)

接下来查找体素间距最大（分辨率最低、target中最大的数）的轴，存入worst_spacing_axis变量；将其余轴一起存入other_axes变量中并获取这些轴体素间距和图像大小的中位数。

worst_spacing_axis = np.argmax(target)
other_axes = [i for i in range(len(target)) if i != worst_spacing_axis]
other_spacings = [target[i] for i in other_axes]
other_sizes = [target_size[i] for i in other_axes]

接下来根据ANISO_THRESHOLD 变量（数值为3）判断数据集是否具有各向异性，nnUNet作者认为当一个轴的体素间距大于其他轴三倍以上，就有各向异性。
ANISO_THRESHOLD 变量在__init__函数中存入self.anisotropy_threshold变量。

has_aniso_spacing = target[worst_spacing_axis] > (self.anisotropy_threshold * max(other_spacings))
has_aniso_voxels = target_size[worst_spacing_axis] * self.anisotropy_threshold < min(other_sizes)

如果该数据集具有各向异性则调整其体素间距：

1. 获取数据集中体素间距最大的轴上的体素间距（spacings），存入spacings_of_that_axis变量

2. 计算该轴（spacings_of_that_axis变量）上的第10百分位数（从小到大排序，第10% * 数组长度 位置的数）的体素间距存入target_spacing_of_that_axis变量

3. 判断该体素间距（target_spacing_of_that_axis变量）是否小于其余轴的体素间距的最大值，小于则将其余轴的体素间距的最大值和target_spacing_of_that_axis变量二者的最大值重新赋值给
   target_spacing_of_that_axis变量，再加上1e-5（保证新的target_spacing_of_that_axis大于其他轴的体素间距）。

4. 最终体素间距最大的轴的体素间距更新为target_spacing_of_that_axis变量。

上面这段代码是将具有各向异性的数据集中体素间距最大的轴的体素间距缩小至合理范围。

if has_aniso_spacing and has_aniso_voxels:
    spacings_of_that_axis = spacings[:, worst_spacing_axis]
    target_spacing_of_that_axis = np.percentile(spacings_of_that_axis, 10)
    # don't let the spacing of that axis get higher than the other axes
    if target_spacing_of_that_axis < max(other_spacings):
        target_spacing_of_that_axis = max(max(other_spacings), target_spacing_of_that_axis) + 1e-5
    target[worst_spacing_axis] = target_spacing_of_that_axis

函数最后返回target变量，该变量存储了调整后的各个轴的体素间距。

三. determine_transpose函数

1. 参数

无具体参数

2. 过程

本函数确定数据各个轴是否转置。
首先判断是否强制不转置（self.suppress_transpose变量，在__init__函数中定义），作者表明这个变量后续改进会用到，当前默认是False

if self.suppress_transpose:
    return [0, 1, 2], [0, 1, 2]

获取各轴上调整后的体素间距中位数，存入target_spacing变量，determine_fullres_target_spacing函数见本文。

接下来找到体素间距最大的轴，以及其余轴，分别存入max_spacing_axis变量和remaining_axes变量，最终确定前向转置（最大的轴在第一位，其余轴按原顺序随后）和后向转置

max_spacing_axis = np.argmax(target_spacing)
remaining_axes = [i for i in list(range(3)) if i != max_spacing_axis]
transpose_forward = [max_spacing_axis] + remaining_axes
transpose_backward = [np.argwhere(np.array(transpose_forward) == i)[0][0] for i in range(3)]

最后返回前向和后向转置数组

四. static_estimate_VRAM_usage函数

1. 参数

• patch_size：输入网络的图像大小
• input_channels：输入通道数
• output_channels：输出通道数
• arch_class_name：神经网络模型名称，类名，例如PlainConvUNet
• arch_kwargs：网络模型相关参数
• arch_kwargs_req_import：网络额外需要import的参数

2. 过程

本函数根据nnUNet同作者创建的PlainConvUNet和ResidualEncoderUNet两类U-Net结构计算出合理的显存使用量，两个类中均包含计算本网络单元显存使用量的类函数。
首先获取当前系统中torch允许使用的线程数，用于函数结束时还原

a = torch.get_num_threads()

设置torch允许使用的线程数为 get_allowed_n_proc_DA() 函数返回的值，该函数会返回不同系统上设置可使用的进程数量。

torch.set_num_threads(get_allowed_n_proc_DA())

接下来使用给定的架构参数实例化网络，计算神经网络的显存使用量，再恢复torch原本的线程数设置。

net = get_network_from_plans(arch_class_name, arch_kwargs, arch_kwargs_req_import, input_channels,
                                output_channels,
                                allow_init=False)
ret = net.compute_conv_feature_map_size(patch_size)
torch.set_num_threads(a)

五. determine_resampling函数

本函数确定并返回image和seg的重采样函数及其对应参数。代码很清晰，就不粘贴了。

六. determine_segmentation_softmax_export_fn函数

本函数确定并返回预测seg的重采样函数及其对应参数，在运行nnUNetv2_predict命令时使用。代码很清晰，就不粘贴了

七. 涉及的函数

1. get_filenames_of_train_images_and_targets函数

见此处

2. _maybe_copy_splits_file函数

本函数将一个json文件复制到目标文件夹下

a. 参数

• splits_file： 分割的配置文件路径
• target_fname： 目标文件路径

b. 过程

如果目标文件不存在，直接复制。
如果存在则通过断言判断两个文件是否相同。
代码较简单，就不粘贴了