复现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_channels
self.output_channels = output_channels
stride = maybe_convert_scalar_to_list(conv_op, stride)
self.stride = stride
kernel_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 x
def 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 padding
return 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)