MedNeXt

版权所有 © 德国癌症研究中心(DKFZ)，医学图像计算部门(MIC)。请确保您对此代码的使用符合代码许可。

MedNeXt是一个专为3D医学图像分割设计的完全ConvNeXt架构，旨在利用ConvNeXt块的可扩展性，同时针对稀疏标注的医学图像分割数据集的挑战进行定制。MedNeXt是一个正在开发中的模型，预计在不久的将来会定期更新。

当前的训练框架建立在nnUNet (v1)之上 - 模块名称nnunet_mednext反映了这一点。您可以自由地将该架构应用于您自己的训练流程，或使用本仓库中的流程。我们为这两种方式都提供了说明。

[arXiv, 2023] [MICCAI 2023]

如果您发现这个模型对您的研究有用，请引用以下工作：

Roy, S., Koehler, G., Ulrich, C., Baumgartner, M., Petersen, J., Isensee, F., Jaeger, P.F. & Maier-Hein, K. (2023).
MedNeXt: Transformer-driven Scaling of ConvNets for Medical Image Segmentation. 
International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), 2023.

如果您使用此流程进行训练，请同时引用以下工作：

Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2020). 
nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nature Methods, 1-9.

目录

• MedNeXt
• 目录
• 当前版本和显著特性
• 安装
• MedNeXt架构及在外部流程中的使用
  • MedNeXt v1
    • 作为完整MedNeXt v1架构使用
    • MedNeXt块的单独使用
    • UpKern权重加载
• 内部训练流程的使用
  • 计划和预处理
  • 使用nnUNet (v1)训练MedNeXt
  • 使用UpKern训练5x5x5内核版本

当前版本和显著特性：

• v1 (MICCAI 2023)：完全3D ConvNeXt架构、残差ConvNeXt重采样、用于大内核的UpKern、用于训练大型模型的梯度检查点

如前所述，MedNeXt正在积极开发中，预计未来版本将对流程进行进一步改进。

安装

可以使用以下命令克隆和安装仓库。

git clone https://github.com/MIC-DKFZ/MedNeXt.git mednext
cd mednext
pip install -e .

MedNeXt架构及在外部流程中的使用

MedNeXt可用于外部3D体积分割训练流程，类似于任何PyTorch nn.Module。在仅作为架构使用时，它在功能上与nnUNet解耦。只需安装仓库并导入架构或块即可。理论上，可以使用MedNeXt作为编码器-解码器风格的网络以及块来自由定制网络。

MedNeXt v1

MedNeXt v1是MedNeXt的第一个版本，包含了此处描述的架构特性。

重要提示： MedNeXt v1是以1.0mm等距间隔进行训练的，这是UNETR、SwinUNETR等架构所偏好的。虽然理论上完全可以使用其他间隔（如原生nnUNet偏好的中位数间隔），但目前尚未在MedNeXt v1上进行测试，可能会影响性能。

下面描述了如何将整个MedNeXt v1作为完整架构使用，以及如何在外部架构中使用MedNeXt块。

作为完整MedNeXt v1架构使用：

可以通过以下方式导入架构，并提供多个参数。

from nnunet_mednext.mednextv1 import MedNeXt

model = MedNeXt(
          in_channels: int,                         # 输入通道数
          n_channels: int,                          # 基础通道数
          n_classes: int,                           # 类别数
          exp_r: int = 4,                           # 扩展层中的扩展比率
          kernel_size: int = 7,                     # 深度卷积层的内核大小
          enc_kernel_size: int = None,              # 编码器中的（单独）内核大小
          dec_kernel_size: int = None,              # 解码器中的（单独）内核大小
          deep_supervision: bool = False,           # 启用深度监督
          do_res: bool = False,                     # MedNeXt块中的残差连接
          do_res_up_down: bool = False,             # 重采样块中的残差连接
          checkpoint_style: bool = None,            # 启用梯度检查点
          block_counts: list = [2,2,2,2,2,2,2,2,2], # 每层的深度优先块数 
          norm_type = 'group',                      # 归一化类型：'group'或'layer'
          dim = '3d'                                # 支持'3d'、'2d'参数
)

请注意，1）深度监督和2）MedNeXt和上/下采样块中的残差连接在发表的论文中都用于训练。

梯度检查点可用于在低内存设备上训练更大的模型，通过计算换取激活存储。此版本中实现的检查点是在MedNeXt块级别。

MedNeXt v1已经测试了4种定义的架构大小和2种定义的内核大小。它们的具体情况如下：

定义了实用函数，用于重新创建这些架构（带或不带深度监督），根据输入通道数、目标类别数、论文中使用的模型ID、内核大小和深度监督进行定制：

from nnunet_mednext import create_mednext_v1

model = create_mednext_v1(
  num_channels = 3,
  num_classes = 10,
  model_id = 'B',             # S、B、M和L是有效的模型ID
  kernel_size = 3,            # 论文中测试了3x3x3和5x5x5
  deep_supervision = True     # 在论文中使用
)

MedNeXt块的单独使用

MedNeXt块可以像整个架构一样单独导入使用。可以直接导入以下块进行使用。

from nnunet_mednext import MedNeXtBlock, MedNeXtDownBlock, MedNeXtUpBlock

# 标准MedNeXt块
block = MedNeXtBlock(
    in_channels:int,              # 输入通道数
    out_channels:int,             # 输出通道数
    exp_r:int=4,                  # 扩展层中的通道扩展比率
    kernel_size:int=7,            # 深度卷积层的内核大小
    do_res:bool=True,              # 是否使用残差连接。默认：True
    norm_type:str = 'group',      # 归一化类型：'group'或'layer'
    n_groups:int or None = None,  # 深度卷积层中的组数
                                  # （在大多数情况下保持为'None'）
)


# 带MedNeXt块的2倍下采样
block_down = MedNeXtDownBlock(
    in_channels:int,              # 输入通道数
    out_channels:int,             # 输出通道数
    exp_r:int=4,                  # 扩展层中的通道扩展比率
    kernel_size:int=7,            # 深度卷积层的内核大小
    do_res:bool=True,              # 是否使用残差连接。默认：True
    norm_type:str = 'group',      # 归一化类型：'group'或'layer'
)


# 带MedNeXt块的2倍上采样
block_up = MedNeXtUpBlock(
    in_channels:int,              # 输入通道数
    out_channels:int,             # 输出通道数
    exp_r:int=4,                  # 扩展层中的通道扩展比率
    kernel_size:int=7,            # 深度卷积层的内核大小
    do_res:bool=True,              # 是否使用残差连接。默认：True
    norm_type:str = 'group',      # 归一化类型：'group'或'layer'
)

UpKern权重加载

UpKern是一种简单的算法，用于使用等效小内核MedNeXt初始化大内核MedNeXt网络。等效是指具有相同配置的网络，唯一的区别是深度卷积层中的内核大小。 大内核通过对其较小对应物进行三线性插值来初始化。 以下是使用此权重加载样式的示例。

from nnunet_mednext import create_mednext_v1
from nnunet_mednext.run.load_weights import upkern_load_weights
m_net_ = create_mednext_v1(1, 3, 'S', 5)
m_pre = create_mednext_v1(1, 3, 'S', 3)

# 通常m_pre是预训练的
m3 = upkern_load_weights(m_net_, m_pre)

内部训练流程的使用

计划和预处理

要按照MICCAI 2023版本预处理数据集，请运行：

mednextv1_plan_and_preprocess -t YOUR_TASK -pl3d ExperimentPlanner3D_v21_customTargetSpacing_1x1x1 -pl2d ExperimentPlanner2D_v21_customTargetSpacing_1x1x1

与nnUNet一样，如果不需要某些维度的预处理数据，可以将-pl3d或-pl2d设置为None。 请注意，本仓库中的YOUR_TASK设计为旧版nnUNet(v1)格式。如果要使用最新的nnUNet(v2)，您需要自行调整预处理器。

自定义的ExperimentPlanner3D_v21_customTargetSpacing_1x1x1旨在将patch大小设置为128x128x128，间距设置为1mm等距，因为这些是MICCAI 2023版本中使用的实验条件。

使用nnUNet(v1)训练方法训练MedNeXt

MedNeXt有自定义的nnUNet(v1)训练器，允许它以类似于基础架构的方式进行训练。 如果您不熟悉这种代码格式，请查看nnUNet仓库中的旧版nnUNet(v1)分支。请查看这里以获取所有可用的训练器，以重现MICCAI 2023实验。请注意，所有训练器都是3D的，因为该架构在3D模式下进行了测试。当然，您也可以创建自定义训练器（包括2D架构的2D训练器）。

mednextv1_train 3d_fullres TRAINER TASK_NUMBER FOLD -p nnUNetPlansv2.1_trgSp_1x1x1

有4种架构（S、B、M、L）和2种核大小（3、5）的训练器，用于复现MICCAI 2023的实验。 以下是在任务Task040_KiTS2019的5折交叉验证中的第0折上训练nnUNetTrainerV2_MedNeXt_S_kernel3训练器的示例：

mednextv1_train 3d_fullres nnUNetTrainerV2_MedNeXt_S_kernel3 Task040_KiTS2019 0 -p nnUNetPlansv2.1_trgSp_1x1x1

也可以通过这种方式从头开始训练5x5x5核版本，但我们建议先训练3x3x3核版本，然后使用UpKern。

使用UpKern训练5x5x5核版本

要使用UpKern训练5x5x5核版本，必须先训练3x3x3核版本。要使用UpKern进行训练，只需运行以下命令：

mednextv1_train 3d_fullres TRAINER TASK FOLD -p nnUNetPlansv2.1_trgSp_1x1x1 -pretrained_weights YOUR_MODEL_CHECKPOINT_FOR_KERNEL_3_FOR_SAME_TASK_AND_FOLD -resample_weights

以下是使用UpKern在任务Task040_KiTS2019的5折交叉验证中的第0折上训练nnUNetTrainerV2_MedNeXt_S_kernel5训练器的示例：

mednextv1_train 3d_fullres nnUNetTrainerV2_MedNeXt_S_kernel5 Task040_KiTS2019 0 -p nnUNetPlansv2.1_trgSp_1x1x1 -pretrained_weights SOME_PATH/nnUNet/3d_fullres/Task040_KiTS2019/nnUNetTrainerV2_MedNeXt_S_kernel3__nnUNetPlansv2.1_trgSp_1x1x1/fold_0/model_final_checkpoint.model -resample_weights

-resample_weights标志负责触发UpKern算法。

关于2D MedNeXt的说明：

请注意，虽然MedNeXt可以在2D模式下运行，但尚未在2D模式下进行测试。