复现nnUNet2并跑通自定义数据

1. 配置环境

stage1：创建python环境，这里建议python3.10

conda create --n nnunet python=3.10
conda activate nnunet 

stage2：按照pytorch官方链接 (conda/pip) 上的说明安装PyTorch。
nnunetv2 2.5.1 需要 torch>=2.1.2，比如：

conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=11.8

但是，使用上述指令安装pytorch，很多时候还是装的cpuonly版，这里可以进行本地安装，首先搜索whl，最后的cuda版本可以自行更改。

https://download.pytorch.org/whl/cu117

最后，直接cd到下载文件所属的盘符下：

d:
cd D:\program
pip install torch-2.1.2+cu118-cp310-cp310-win_amd64.whl

stage3：git nnUnet的源代码，并配置nnUnet

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

2. 处理数据集

可以通过运行将 nnU-Net v1 中的数据集转换为 V2 nnUNetv2_convert_old_nnUNet_dataset INPUT_FOLDER OUTPUT_DATASET_NAME。请记住，v2 调用数据集 DatasetXXX_Name（而不是 Task），其中 XXX 是一个 3 位数字。请提供旧任务的路径，而不仅仅是任务名称。nnU-Net V2 不知道 v1 任务在哪里！

2.1 创建文件夹

首先，创建nnUNet_raw、nnUNet_preprocessed、nnUNet_results三个文件夹，将下载的原始数据集放到nnUNet_raw下，这里MSD挑战赛的Task02_Heart数据集为例。
nnUNet_raw用于存放原始数据集；
nnUNet_preprocessed用于存放格式转换并预处理的数据集及配置信息；
nnUNet_results用于存放训练结果相关的信息文件。

2.2 数据集格式转换

接着，我们转换数据集格式：

nnUNetv2_convert_MSD_dataset -i D:\nnUNet\nnUNet_raw\Task02_Heart

可能会出现下列错误：

TypeError: expected str, bytes or os.PathLike object, not NoneType

解决方法：
原因是没有设置相关数据集的路径，打开D:\nnUNet\nnunetv2\paths.py文件

方法1：直接在这个py文件指定数据集路径：

方法2：
设置环境变量：
1.windows：

#这个变量只在当前命令提示符会话中有效，不会被保存到系统环境变量中
set nnUNet_raw=D:/nnUNet/nnUNet_raw
set nnUNet_preprocessed=D:/nnUNet/nnUNet_preprocessed
set nnUNet_results=D:/nnUNet/nnUNet_results

2.linux：

export nnUNet_raw="D:/nnUNet/nnUNet_raw"
export nnUNet_preprocessed="D:/nnUNet/nnUNet_preprocessed"
export nnUNet_results="D:/nnUNet/nnUNet_results"

发现，转换成功，生成了1个Dataset002_Heart文件夹：

转换前后主要区别在于：
1、文件夹由Task02_Heart变成Dataset002_Heart。
ps：文件夹命名为：Dataset+三位整数+任务名，Dataset002_Heart中数据集ID为2，任务名为Heart。此文件夹下存放需要的训练数据集imageTr、测试集imageTs、标签labelsTr。其中labelsTr是与imageTr中一一对应的标签，文件中都是nii.gz文件。imageTs是可选项，可以没有。
2、原始json中是：

转换后，字典的key为"channel_names"：

3、原始json中有"training"、"test"这两个key存放数据集的相对路径：

转换后，字典的key为"file_ending"，直接记录数据集的后缀名即可：

2.3 数据集预处理

nnUNetv2_plan_and_preprocess -d DATASET\_ID --verify_dataset_integrity  
#这里我是用的是Dataset002_Heart，将ID改为2
nnUNetv2_plan_and_preprocess -d 2 --verify_dataset_integrity

运行成功后，会在nnUNet_preprocessed文件下生成如下文件：

此步骤主要就是对数据进行：裁剪crop，重采样resample以及标准化normalization。将提取数据集fingerprint（一组特定于数据集的属性，例如图像大小、体素间距、强度信息等）。此信息用于设计三种 U-Net 配置。每个管道都在其自己的数据集预处理版本上运行。其中的nnUNetPlans.json是网络结构相关的配置文件。

3. 训练

训练使用如下脚本指令：

nnUNetv2_train DATASET_NAME_OR_ID UNET_CONFIGURATION FOLD [additional options, see -h]

DATASET_NAME_OR_ID 指定应在哪个数据集上进行训练；UNET_CONFIGURATION 是一个字符串，用于标识所请求的 U-Net 配置（默认值：2d、3d_fullres、3d_lowres、3d_cascade_lowres）；
FOLD 指定训练 5 折交叉验证中的哪一折。
根据数据集名称或者数据集的ID、UNET配置、折数FOLD，进行训练：

nnUNetv2_train 2 2d 0

成功开始训练：

4. 改进模型

4.1 概要

许多人可能不满足于只是复现跑通，而是有一些改进模型，提升指标，最后发paper的需求。
但是，会发现nnUNet2的网络模型结构很难找，但当找到nnUNet\nnunetv2\utilities\get_network_from_plans.py时，发现还是找不到网络模型的model。
但是，我们在数据集处理的时候生成的nnUNet_preprocessed文件下的nnUNetPlans.json文件中发现了网络结构的model文件。

于是，发现这个模型结构在我们配置的环境中，我就在Anaconda中我的环境下D:\Anaconda3\envs\nnunet\Lib\site-packages\dynamic_network_architectures\architectures找到了。

4.2 加注意力模块

我们以在unet的编码器中添加SE注意力模块为例，首先打开unet.py文件，找到调用编码器实例化类对象的地方。

发现，PlainConvEncoder类中卷积模块通过StackedConvBlocks类实例化出来的，

简单点，我们就直接在StackedConvBlocks类中嵌入SE模块。找到了，nn.Sequential()模块拼接模块的部分，
我们在下面堆上SE模块。

这里我也附上SE的代码，通过参考其它来的继承来写SE：

class SE(nn.Module):def __init__(self,conv_op: Type[_ConvNd],input_channels: int,output_channels: int,kernel_size: Union[int, List[int], Tuple[int, ...]],stride: Union[int, List[int], Tuple[int, ...]],conv_bias: bool = False,norm_op: Union[None, Type[nn.Module]] = None,norm_op_kwargs: dict = None,dropout_op: Union[None, Type[_DropoutNd]] = None,dropout_op_kwargs: dict = None,nonlin: Union[None, Type[torch.nn.Module]] = None,nonlin_kwargs: dict = None,nonlin_first: bool = False):super(SE, self).__init__()self.input_channels = input_channelsself.output_channels = output_channelsstride = maybe_convert_scalar_to_list(conv_op, stride)self.stride = stridekernel_size = maybe_convert_scalar_to_list(conv_op, kernel_size)if norm_op_kwargs is None:norm_op_kwargs = {}if nonlin_kwargs is None:nonlin_kwargs = {}ops = []self.conv = conv_op(input_channels,output_channels,kernel_size,stride,padding=[(i - 1) // 2 for i in kernel_size],dilation=1,bias=conv_bias,)ops.append(self.conv)if dropout_op is not None:self.dropout = dropout_op(**dropout_op_kwargs)ops.append(self.dropout)if norm_op is not None:self.norm = norm_op(output_channels, **norm_op_kwargs)ops.append(self.norm)self.all_modules = nn.Sequential(*ops)self.se = SEAttention(conv_op=conv_op, channel=output_channels)def forward(self, x):x = self.all_modules(x)x = self.se(x)return xdef compute_conv_feature_map_size(self, input_size):assert len(input_size) == len(self.stride), "just give the image size without color/feature channels or " \"batch channel. Do not give input_size=(b, c, x, y(, z)). " \"Give input_size=(x, y(, z))!"output_size = [i // j for i, j in zip(input_size, self.stride)]  # we always do same paddingreturn np.prod([self.output_channels, *output_size], dtype=np.int64)class SEAttention(nn.Module):def __init__(self, conv_op, channel=512, reduction=16):super().__init__()if conv_op == torch.nn.modules.conv.Conv2d:self.pool = nn.AdaptiveAvgPool2d(1)elif conv_op == torch.nn.modules.conv.Conv3d:self.pool = nn.AdaptiveAvgPool3d(1)self.fc = nn.Sequential(nn.Linear(channel, channel // reduction, bias=False),nn.ReLU(inplace=True),nn.Linear(channel // reduction, channel, bias=False),nn.Sigmoid())def forward(self, x):size = x.size()b, c = size[0], size[1]y = self.pool(x).view(b, c)y = self.fc(y).view(b, c, *[1 for i in range(len(size) - 2)])return x * y.expand_as(x)def init_weights(self):for m in self.modules():if isinstance(m, nn.Linear):init.normal_(m.weight, std=0.001)if m.bias is not None:init.constant_(m.bias, 0)