nnUNet使用笔记：如何自定义你的nnUNet？

简介

nnUNet在许多医学challenge的分割任务上都达到了SOTA水平，身边一些人也实际使用过，验证了有效性。nnUNet的一大优点就是可以直接命令行调用，无需手动调参；但这也是他的问题，手动修改网络/Loss/超参的方式并不是那么简明。

本文主要是笔者的个人使用笔记，用于记录如何调整nnUNet框架下的lr、epoch、网络结构、loss函数等配置。内容主要是分享如何在享受nnUNet的方便性的同时，也能够享受调参的快 (zhe) 乐 (mo)。

正文

1. 修改lr，epoch

这个可以简单的通过修改Trainer源码来实现，比如对于

nnUNet_train [model] nnUNetTrainerV2 [task] [fold] --fp32

可以修改nnUNetTrainerV2.py对应的源码。

class nnUNetTrainerV2(nnUNetTrainer):
    """
    Info for Fabian: same as internal nnUNetTrainerV2_2
    """

    def __init__(self, plans_file, fold, output_folder=None, dataset_directory=None, batch_dice=True, stage=None,
                 unpack_data=True, deterministic=True, fp16=False):
        super().__init__(plans_file, fold, output_folder, dataset_directory, batch_dice, stage, unpack_data,
                         deterministic, fp16)
        self.max_num_epochs = config.epoch
        self.initial_lr = config.lr
        self.deep_supervision_scales = None
        self.ds_loss_weights = None

        self.pin_memory = True

2. 修改loss函数、网络模型（改用UNet变种或者其他）

默认loss是nnunet/training/loss_functions/dice_loss.py中的DC_and_CE_loss，一种不太便于维护的方法就是，直接修改这个Loss的源码，在准备接下来的实验中都使用同一个Loss时可以这么搞。

比较建议的方式是，通过修改plan中的ExperimentPlanner，然后重新进行nnUNet_plan_and_preprocess来修改(网络模型也同样可以使用这种方式进行修改)。

[TO-DO] 给个例子说明