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

w1ndfly

已于 2025-02-27 23:04:55 修改

阅读量734

点赞数 23

分类专栏：阅读nnUNet V2代码文章标签：人工智能机器学习深度学习计算机视觉卷积神经网络

于 2025-01-10 11:43:27 首次发布

本文链接：https://blog.youkuaiyun.com/shey_joe/article/details/145048937

版权

阅读nnUNet V2代码专栏收录该内容

10 篇文章

订阅专栏

前文请见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（三）

plan_experiment函数

参数

无实际参数

过程

本函数主要配置预处理、训练、后处理过程中用到的常量数据、处理函数、网络结构等信息，并将其保存为nnUNetPlans.json文件。

创建一个缓存字典（_tmp变量）；再获取前向转置和后向转置（由determine_transpose函数获得）；之后获取全分辨率的体素间距（由determine_fullres_target_spacing获得）并前向转置它；代码较为清晰，不做粘贴。

接下来依据dataset_fingerprint.json文件（存储在self.dataset_fingerprint变量）内的spacings和shapes_after_crop以及上文得到的全分辨率体素间距（fullres_spacing）三个变量计算数据集内图像的新形状；计算出新shape后，找到其三个维度的中值，并将其前向转置，得到前向转置的、合理优化的、中位数的数据集image和seg的shape，存入new_median_shape_transposed变量，这只是参考值，后续代码有优化。

new_shapes = [compute_new_shape(j, i, fullres_spacing) for i, j in
                zip(self.dataset_fingerprint['spacings'], self.dataset_fingerprint['shapes_after_crop'])]
new_median_shape = np.median(new_shapes, 0)
new_median_shape_transposed = new_median_shape[transpose_forward]

其中，compute_new_shape函数计算方式为 (shapes_after_crop * spacings) / fullres_spacing。

接下来根据刚刚得到的shape（new_median_shape_transposed变量）和训练集样例数量计算（估算）出整个训练集的image体素数量：

approximate_n_voxels_dataset = float(np.prod(new_median_shape_transposed, dtype=np.float64) *
                               self.dataset_json['numTraining'])

前菜准备完毕，开始配置2d、3d_fullres和3d_lowres

本函数针对2d、3d_fullres和3d_lowres三种情况提供了相应的训练配置，配置详见get_plans_for_configuration函数。如果只有2d则不配置另外两个3d：

if new_median_shape_transposed[0] != 1:	## 3d情况
	# 省略**3d_fullres**和**3d_lowres**代码 ！！！！！
else:
    plan_3d_fullres = None
    plan_3d_lowres = None

首先是3d_fullres和3d_lowres，二者是一起的，先配置3d_fullres，后根据3d_fullres配置3d_lowres。

此处会涉及到3D级联（3D U-Net Cascade），做个说明，便于理解后面配置3d_lowres，见下图请添加图片描述
📌📌在用户使用3D级联训练网络时，nnUNet V2会生成两个U-Net，第一个U-Net输入输出图像是原有数据样例下采样后得到的，根据这个下采样的数据样例（分辨率变低了）大小配置第一个U-Net结构（配置过程详见get_plans_for_configuration函数）。对于第一个U-Net输出的seg图像，nnUNet V2会将其上采样至原有大小，再将上采样后的seg图像按类别（通道）进行one-hot编码后，作为第二个U-Net的数据样例。上图stage 2中的输入输出图像大小、网络结构和stage 1的很像，这并不是因为下采样，而是使用了patch训练的缘故。

说明完毕

3d_fullres的配置简单、直接、一次成型：调用get_plans_for_configuration函数获取对应配置。👍👍：

plan_3d_fullres = self.get_plans_for_configuration(fullres_spacing_transposed,
                                                    new_median_shape_transposed,
                                                    self.generate_data_identifier('3d_fullres'),
                                                    approximate_n_voxels_dataset, _tmp)

配置3d_fullres后，会试着配置3d_lowres，如果不满足要求（即本次数据集不能使用3D级联），就不再配置本数据集的3d_lowres。💔💔

如何判断不能使用3D级联呢？结合上面的3D级联来说明，对于stage 1阶段下采样后的数据样例，如果它的大小 * 2后小于3d_fullres的数据样例大小，就使用，否则，不使用。一句话，二者图像大小差别大就使用，差别小，用了没效果，就不使用。

开始试着配置3d_lowres：获得3d_fullres配置下的patch size，以及上文计算出的shape（new_median_shape_transposed变量，原有数据样例大小的中位数），由这两个变量分别计算一个patch内的体素数量（num_voxels_in_patch变量）和一个样例内的体素数量（median_num_voxels）。接着初始化plan_3d_lowres为None，同时复制一份3d_fullres的体素间距给3d_lowres，再设置间距增加因子（spacing_increase_factor 变量），之前是1.01（见注释），现在是1.03：

patch_size_fullres = plan_3d_fullres['patch_size']
median_num_voxels = np.prod(new_median_shape_transposed, dtype=np.float64)
num_voxels_in_patch = np.prod(patch_size_fullres, dtype=np.float64)

plan_3d_lowres = None
lowres_spacing = deepcopy(plan_3d_fullres['spacing'])spacing_increase_factor = 1.03  # used to be 1.01 but that is slow with new GPU memory estimation!

接下来试着增加3d_lowres配置下的spacing（此时3d_lowres的spacing和3d_fullres一样）：

创建while循环，循环条件是一个patch中的体素数量与一个样例内的体素数量的比值小于设置的阈值（self.lowres_creation_threshold变量，值为0.25）。结束时说明3D级联的stage 1的U-Net的patch size与它的样例大小保持在阈值之内，不让patch size 太小。既要让stage 1的U-Net的样例大小小于原有样例大小，又要让stage 1的U-Net的patch size不能太小，二者缺一不可，共同决定是否能使用3D级联

while循环内代码先获取当前 3d_lowres体素间距（lowres_spacing变量）的最大值，如果该轴间距与其他间距比值大于2，则让其他轴的间距乘间距增加因子，否则所有轴间距都乘间距增加因子。

再基于最新的3d_lowres体素间距、之间获取的3d_fullres体素间距和样例大小的中位数，计算新的、3d_lowres样例内的体素数量（median_num_voxels变量）：

max_spacing = max(lowres_spacing)
if np.any((max_spacing / lowres_spacing) > 2):
    lowres_spacing[(max_spacing / lowres_spacing) > 2] *= spacing_increase_factor
else:
    lowres_spacing *= spacing_increase_factor
median_num_voxels = np.prod(plan_3d_fullres['spacing'] / lowres_spacing * new_median_shape_transposed,
                            dtype=np.float64)

接着通过 self.get_plans_for_configuration 函数获取3d_lowres配置，并根据新的配置计算3d_lowres配置的patch内体素数量（num_voxels_in_patch变量）：

plan_3d_lowres = self.get_plans_for_configuration(lowres_spacing,
                                                    tuple([round(i) for i in plan_3d_fullres['spacing'] /
                                                            lowres_spacing * new_median_shape_transposed]),
                                                    self.generate_data_identifier('3d_lowres'),
                                                    float(np.prod(median_num_voxels) *
                                                        self.dataset_json['numTraining']), _tmp)
num_voxels_in_patch = np.prod(plan_3d_lowres['patch_size'], dtype=np.int64)

while循环结束

如果3d_fullres样例内体素数量与3d_lowres样例内体素数量比值小于2，则丢弃3d_lowres配置，并打印相应信息，设置plan_3d_lowres为None，代码较为清晰，不做粘贴。

接着依据3d_lowres配置是否为None，设置训练时损失函数是否使用batch_dice（batch_dice在训练阶段的损失函数处会介绍）：

if plan_3d_lowres is not None:
    plan_3d_lowres['batch_dice'] = False
    plan_3d_fullres['batch_dice'] = True
else:
    plan_3d_fullres['batch_dice'] = False

处理完3d_fullres和3d_lowres配置后，开始处理2d配置：通过 self.get_plans_for_configuration 函数获取2d配置，并设置batch_dice为True，简单、清晰、一步到位👍👍：

plan_2d = self.get_plans_for_configuration(fullres_spacing_transposed[1:],
                                            new_median_shape_transposed[1:],
                                            self.generate_data_identifier('2d'), approximate_n_voxels_dataset,
                                            _tmp)

接下来获取原始体素间距的中位值以及image前景区域大小的中位值，并对二者前向转置，后续存入nnUNetPlans.json文件中；再将用户设置的dataset.json文件复制到nnUNet_preprocessed对应数据集文件夹下；再将上述获取的所有配置及其他信息存入plan字典中，将其保存为nnUNet_preprocessed对应数据集文件夹下的nnUNetPlans.json。

接下来如果有3d_fullres或3d_lowres配置，就在plan字典中加入对应配置：

if plan_3d_lowres is not None:
    plans['configurations']['3d_lowres'] = plan_3d_lowres
    if plan_3d_fullres is not None:
        plans['configurations']['3d_lowres']['next_stage'] = '3d_cascade_fullres'
    print('3D lowres U-Net configuration:')
    print(plan_3d_lowres)
    print()
if plan_3d_fullres is not None:
    plans['configurations']['3d_fullres'] = plan_3d_fullres
    print('3D fullres U-Net configuration:')
    print(plan_3d_fullres)
    print()
    if plan_3d_lowres is not None:
        plans['configurations']['3d_cascade_fullres'] = {
            'inherits_from': '3d_fullres',
            'previous_stage': '3d_lowres'
        }

📌📌 3d_lowres是3D级联所用，必须搭配3d_cascade_fullres一起，3d_cascade_fullres配置在使用时会照搬3d_fullres配置，但nnUNet V2为了区分当前训练是只有3d_fullres还是3D级联下的“3d_fullres”，加入3d_cascade_fullres字段。

继承（inherits）这个设置还能让用户自定义配置，继承自谁，额外注意，这里的继承是父类覆盖子类，与python的子类覆盖父类不一样。

最后，保存plan字典，并返回它：

self.plans = plans
self.save_plans(plans)
return plans