简介
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] 给个例子说明