MedNeXt

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.

当前版本和显著特性：

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）	内核大小	参数数	GFlops
Small (S)	3x3x3	5.6M	130
Small (S)	5x5x5	5.9M	169
Base (B)	3x3x3	10.5M	170
Base (B)	5x5x5	11.0M	208
Medium (M)	3x3x3	17.6M	248
Medium (M)	5x5x5	18.3M	308
Large (L)	3x3x3	61.8M	500
Large (L)	5x5x5	63.0M	564

定义了实用函数，用于重新创建这些架构（带或不带深度监督），根据输入通道数、目标类别数、论文中使用的模型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算法。