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（三）

一. determine_reader_writer函数

1. 参数

无实际参数

2. 过程

首先从dataset中选取一个（这里是第一个）数据，以此确定用于读取该文件的类

example_image = self.dataset[self.dataset.keys().__iter__().__next__()]['images'][0]

最后通过determine_reader_writer_from_dataset_json函数确定该类

二. determine_normalization_scheme_and_whether_mask_is_used_for_norm函数

1. 参数

无实际参数

2. 过程

本函数是确定样例选用哪种归一化，以及是否在归一化时使用mask

首先确保用户在nnUNet_raw文件夹下的dataset.json文件中定义了channel_names，这个字段配置数据样例的采集方式，例如CT、MRI等。

在配置nnUNet_raw \ Dataset001_XXXX \ imagesTr以及imagesTs文件夹时，单个样例名称最后的_0000或者_0001用于区分数据的采集方式，配合channel_names字段可以对不同采集方式的样例应用不同的归一化方法（由get_normalization_scheme函数查询字典得到，例如CT使用CTNormalization，默认使用ZScoreNormalization）。代码较清晰，不做粘贴。

modalities = self.dataset_json['channel_names'] if 'channel_names' in self.dataset_json.keys() else \
            self.dataset_json['modality']
normalization_schemes = [get_normalization_scheme(m) for m in modalities.values()]

接下来判断dataset_fingerprint.json文件内median_relative_size_after_cropping值是否小于0.75（裁剪image非零区域后图像大小/裁剪前大小），小于说明整个数据集裁剪后的图像显著小于裁剪前的图像，此时根据归一化方法的特性确定是否在归一化时将mask外的像素保持为零（就是不考虑mask外的区域）。不小于则对所有数据样例设置为False。将这些情况存入use_nonzero_mask_for_norm变量。一句话，裁剪太多时，某些归一化方法会不考虑数据样例的边缘零值。

if self.dataset_fingerprint['median_relative_size_after_cropping'] < (3 / 4.):
    use_nonzero_mask_for_norm = [i.leaves_pixels_outside_mask_at_zero_if_use_mask_for_norm_is_true for i in
                                    normalization_schemes]
else:
    use_nonzero_mask_for_norm = [False] * len(normalization_schemes)

最后获取归一化方式的类名，和use_nonzero_mask_for_norm变量一起返回上层函数

三. get_plans_for_configuration函数

1. 参数

• spacing：体素间距

• median_shape：图像大小的中值列表

• data_identifier：标识3d_fullres、2d等

• approximate_n_voxels_dataset：体素数量

• _cache：缓存，防止占用过多计算资源

2. 过程

本函数是确定data_identifier下的训练配置

根据dataset.json文件内设置的channel_names数量确定神经网络输入通道数，用于后续计算大致显存使用量；根据spacing的长度确定最大特征图通道数（其实是查表）；根据spacing长度确定使用nn.Conv2d还是nn.Conv3d（其实也是查表）。

接下来根据图像的体素间距计算初始Patch大小，存入initial_patch_size变量，此处会使初始Patch
Size较大，后续会有优化：

tmp = 1 / np.array(spacing)
if len(spacing) == 3:
    initial_patch_size = [round(i) for i in tmp * (256 ** 3 / np.prod(tmp)) ** (1 / 3)]
elif len(spacing) == 2:
    initial_patch_size = [round(i) for i in tmp * (2048 ** 2 / np.prod(tmp)) ** (1 / 2)]
else:
    raise RuntimeError()

接下来结合median_shape，优化一次initial_patch_size，避免其过大：

initial_patch_size = np.array([min(i, j) for i, j in zip(initial_patch_size, median_shape[:len(spacing)])])

接下来计算U-Net网络结构的相关信息，例如卷积核、池化核等，get_pool_and_conv_props函数见文末：

network_num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, patch_size, \
shape_must_be_divisible_by = get_pool_and_conv_props(spacing, initial_patch_size,
                                                        self.UNet_featuremap_min_edge_length,
                                                        999999)

接下来根据池化核数量（pool_op_kernel_sizes）确定U-Net网络结构会有几个下采样和上采样（num_stages），再根据卷积维度（unet_conv_op）确定实例归一化类型（由get_matching_instancenorm函数获取，依旧是查表），例如unet_conv_op是nn.Conv2d，那么norm就是nn.InstanceNorm2d，为什么使用InstanceNorm？这是nnU-Net论文中要求的 😃：

num_stages = len(pool_op_kernel_sizes)
norm = get_matching_instancenorm(unet_conv_op)

接下来创建字典（architecture_kwargs变量），配置U-Net网络结构，例如使用的U-Net类、下采样层数、每层的特征图数量、卷积维度、kernel_size、stride、编码器和解码器每层的卷积数、卷积是否有bias、归一化操作、非线性激活函数等，这需要结合nnUNet同作者写的dynamic-network-architectures代码（在GitHub上有），了解最终构建的U-Net网络结构 😃。此处代码较为清晰，不做粘贴。

接下来估算显存使用量（estimate变量），详见static_estimate_VRAM_usage函数；计算参考显存值（reference）和参考批次值（ref_bs）

接下来是while循环，期间更新网络结构信息，直至当前网络结构的显存使用量（estimate）除以参考批次大小（ref_bs）再乘以2（训练过程中，保证显存至少能存储两个批次）小于或等于参考显存值（reference）；训练过程中，考虑到数据加载、中间变量存储、多进程处理等因素，需要留出余量，因此要求显存至少要能存储两个批次；

while (estimate / ref_bs * 2) > reference:

while循环开始：
计算当前patch size与median_shape的比例，找出最大比例的轴，该轴（存入axis_to_be_reduced 变量）就是当前需要减小patch size的轴

axis_to_be_reduced = np.argsort([i / j for i, j in zip(patch_size, median_shape[:len(spacing)])])[-1]

接下来是如何减少patch size，结合注释理解：

# 注释机翻：我们不能简单地通过 shape_must_be_divisible_by[axis_to_be_reduced] 来减小该轴，因为这
# 可能会导致我们跳过一些有效的尺寸。例如，对于形状为 256 的情况，shape_must_be_divisible_by 是 64。
# 如果我们减去 64，我们将得到 192，跳过了 224 这个也是有效的补丁大小
# (224 / 2**5 = 7; 7 < 2 * self.UNet_featuremap_min_edge_length(4) 因此它是有效的）。所以我们需要首先
# 减去 shape_must_be_divisible_by，然后重新计算它，然后再减去
# 重新计算的 shape_must_be_divisible_by。很麻烦。
patch_size = list(patch_size)
tmp = deepcopy(patch_size)
tmp[axis_to_be_reduced] -= shape_must_be_divisible_by[axis_to_be_reduced]
_, _, _, _, shape_must_be_divisible_by = \
    get_pool_and_conv_props(spacing, tmp,
                            self.UNet_featuremap_min_edge_length,
                            999999)
patch_size[axis_to_be_reduced] -= shape_must_be_divisible_by[axis_to_be_reduced]

成功减少后，重新计算并更新网络参数（architecture_kwargs）、计算显存使用量（estimate），流程和初始计算一样，代码不再粘贴。
while循环结束
此时网络结构确定，可以确定batch_size 了，但要保证batch_size大于2，小于整个数据集的 5% （防止过拟合）：

batch_size = round((reference / estimate) * ref_bs)
bs_corresponding_to_5_percent = round(
    approximate_n_voxels_dataset * self.max_dataset_covered / np.prod(patch_size, dtype=np.float64))
batch_size = max(min(batch_size, bs_corresponding_to_5_percent), self.UNet_min_batch_size)

根据类内其他函数确定resampling和normalization方案，将上述所有相关变量存入字典（plan）中，并返回

涉及的函数

1. get_pool_and_conv_props函数

参数

• spacing：体素间距
• patch_size：patch_size 😃
• min_feature_map_size：神经网络中，特征图的最小尺寸，一般是4
• max_numpool：池化层最大数量=池化最多次数，初始时很大，999999之类的

过程

首先是初始化必需的变量：获取image维度，也是spacing数组长度；获取spacing和patch_size的备份，之后要循环更新；初始化池化层kernel size（就是1）；初始化卷积层kernel size；初始化各轴池化次数，这个变量在生成nnUNetPlans.json文件中并没有使用；初始化各轴kernel size。

接下来是while循环，不断尝试在各个维度上进行池化操作，直到无法再进行池化（即某个轴的特征图大小小于最小特征图大小的2倍，或者池化次数达到最大限制，或者各轴体素间距spacing差异小于2倍，或者在只剩余一个轴可池化时，该轴大小小于最小特征图大小的3倍）。在每次池化操作中，会根据当前的体素间距和patch size来确定哪些轴可以进行池化，并更新相应的池化kernel size和卷积kernel size。
最后返回调整后的patch size和池化、卷积层的配置信息。