目录
一、简介
nnUnet是有监督的医学图像分割绕不开的话题（虽然看到有些文章对比实验没加nnUnet？可能比不过？ ），其卓越的性能和简易的方法，为相关研究者提供了一项强有力的工具。然而，由于高度封装性，在原先代码中嵌入自定义网络进行训练，并不是十分方便（至少对我来说 ）。本文旨在分享一点在使用nnUnet训练自定义网络过程中的一点经验，可能存在纰漏（？？？），欢迎在讨论区交流！
二、准备工作
2.1 硬件需求
nnUnet的建议环境是Linux，若使用Windows，需修改路径相关代码（斜杠和反斜杠的替换），很麻烦（不推荐）。博主是在Ubuntu环境中使用Pycharm进行nnUnet的学习。
2.2 调试环境
nnUnet官方推荐的使用方法是在命令行，但这不方便初学者学习。因为只用过Pycharm调试代码（菜！），所以为了满足自己的需求（？），以便于用Pycharm的傻瓜式调试按钮，修改了部分代码： nnunet/paths.py 和 nnunet/run/run_training.py
2.2.1 路径
位于***nnunet/paths.py***文件中，将三个变量路径修改为自己的路径。custom_是博主自己定义的文件，大家可以随意实现。
from custom_ import custom_config
base = custom_config['base']
preprocessing_output_dir = custom_config['preprocessing_output_dir']
network_training_output_dir_base = custom_config['network_training_output_dir_base']
2.2.2 parser
位于***nnunet/run/run_training.py***文件中，这里nnUnet训练代码的入口(!!!)。由于不是命令行调用方式，需要将parser进行修改，添加 “-” 并设置 default 值。
parser = argparse.ArgumentParser()
parser.add_argument("-network", default='2d')
parser.add_argument("-network_trainer", default='nnUNetTrainerV2')
parser.add_argument("-task", default='666', help="can be task name or task id")
parser.add_argument("-fold", default='0', help='0, 1, ..., 5 or \'all\'')
三、训练
3.1 构建网络
nnUnet要求网络继承 SegmentationNetwork 类，这里提供一种可实现的方法，用的时候将 self.model 修改为 自定义网络 即可。
from nnunet.network_architecture.neural_network import SegmentationNetwork
class custom_net(SegmentationNetwork):
def __init__(self, num_classes):
super(custom_net, self).__init__()
self.params = {
'content': None}
self.conv_op = nn.Conv2d
self.do_ds = True
self.num_classes = num_classes
######## self.model 设置自定义网络 by Sleeep ########
self.model = None
######## self.model 设置自定义网络 by Sleeep ########
self.name = self.model.name
def forward(self, x):
if self.do_ds:
return [self.model(x), ]
else:
return self.model(x)
def create_model():
return custom_net(num_classes=2)
3.2 修改配置
构建好网络后，还需要修改一些超参数才能完成训练，修改内容位于 /nnunet/training/network_training/nnUNetTrainerV2.py 文件中。修改nnUNetTrainerV2类中的两个函数 initialize 和 initialize_network。为了减少训练代数，还可修改函数***init***
3.2.1 initialize
def initialize(self, training=True, force_load_plans=False):
"""
- replaced get_default_augmentation with get_moreDA_augmentation
- enforce to only run this code once
- loss function wrapper for deep supervision
:param training:
:param force_load_plans:
:return:
"""
if not self.was_initialized:
maybe_mkdir_p(self.output_folder)
if force_load_plans or (self.plans is None):
self.load_plans_file()
# load plan informantion !!! modify batch_size or patch_size after this
self.process_plans(self.plans)
############## modify para by Sleeep ##############
self.patch_size = np.array(custom_config['patch_size']).astype(int)
self.batch_size = custom_config['batch_size']
self.net_num_pool_op_kernel_sizes = [[2, 2]]
############## modify para by Sleeep ##############
self.process_plans(self.plans)： 官方函数，会载入预处理阶段所生成的各种参数
self.patch_size: 官方预处理后的 图像尺寸不一定满足自定义网络的需要，可在这里修改。举个例子，在预处理阶段，nnUnet自动确定的patch_size为[53, 64]（对于2d网络），然而我的网络需要满足输入尺寸均为 32 的整数倍，自动生成的patch_size并不能满足，所以这里可修改为[64, 64]。self.patch_size会在后续构建 数据增强方法 的函数中使用。
self.batch_size: 根据自己的硬件配置修改
self.net_num_pool_op_kernel_sizes： 这个参数非常重要！其作用是 确定 深监督的层数 和 不同层数的尺寸大小。这里默认自定义网络不使用深监督，所以设置 为 只有一个列表元素即可，里面的值随意（可能吧？对于不是使用深监督的情况）
3.2.2 initialize_network
def initialize_network(self):
"""
- momentum 0.99
- SGD instead of Adam
- self.lr_scheduler = None because we do poly_lr
- deep supervision = True
- i am sure I forgot something here
Known issue: forgot to set neg_slope=0 in InitWeights_He; should not make a difference though
:return:
"""
# if self.threeD:
# conv_op = nn.Conv3d
# dropout_op = nn.Dropout3d
# norm_op = nn.InstanceNorm3d
#
# else:
# conv_op = nn.Conv2d
# dropout_op = nn.Dropout2d
# norm_op = nn.InstanceNorm2d
#
# norm_op_kwargs = {'eps': 1e-5, 'affine': True}
# dropout_op_kwargs = {'p': 0, 'inplace': True}
# net_nonlin = nn.LeakyReLU
# net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
# self.network = Generic_UNet(self.num_input_channels, self.base_num_features, self.num_classes,
# len(self.net_num_pool_op_kernel_sizes),
# self.conv_per_stage, 2, conv_op, norm_op, norm_op_kwargs, dropout_op,
# dropout_op_kwargs,
# net_nonlin, net_nonlin_kwargs, True, False, lambda x: x, InitWeights_He(1e-2),
# self.net_num_pool_op_kernel_sizes, self.net_conv_kernel_sizes, False, True, True)
############## add custom model by Sleeep ##############
self.network = create_model()
############## add custom model by Sleeep ##############
if torch.cuda.is_available():
self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
将原先nnunet网络注释，加入自己的网络。
3.2.3 init
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)
##### by Sleeep ####
self.max_num_epochs = custom_config['epoch']
self.initial_lr = custom_config['lr']
##### by Sleeep ####
self.deep_supervision_scales = None
self.ds_loss_weights = None
self.pin_memory = True
默认的epoch是1000，有点久，改小点。默认lr是0.01