一、简介
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网络注释,加入自己的网络。文章来源:https://www.toymoban.com/news/detail-466895.html
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文章来源地址https://www.toymoban.com/news/detail-466895.html
四、注意事项
- 点赞 收藏 评论 再走
- 先过一遍nnunet论文和官方github中的说明。官方提供的一个example ,其中使用的数据集相对来说较小,可先在该数据集上调试成功后,再使用自定义网络
参考资料
- nnUnet
到了这里,关于nnUnet: 使用自己(自定义)网络训练的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!