nnUNet V2修改网络——加入GAB模块

​更换前，要用nnUNet V2跑通所用数据集，证明nnUNet V2、数据集、运行环境等没有问题

阅读nnU-Net V2 的 U-Net结构，初步了解要修改的网络，知己知彼，修改起来才能游刃有余。

EGE-UNet 是 UNet 的一个变体，专为皮肤病变分割而设计。它在 UNet 的基础上引入了 GHPA 模块和 GAB 模块，分别用于多维度特征提取和多尺度信息融合。GHPA 运用多轴分组和 Hadamard 乘积机制，能从不同视角提取病理信息，而 GAB 则通过分组聚合和卷积操作，有效整合不同层次和尺度的特征。这种创新设计使得 EGE-UNet 不仅在 ISIC2017 和 ISIC2018 数据集上取得卓越的分割效果，还大幅降低了参数和计算成本，成为低资源环境下理想的分割模型。

EGE-UNet官方代码仓库： https://github.com/JCruan519/EGE-UNet

GAB模块巧妙地整合了来自不同深度的高低层特征。它先对特征图进行通道分组，将低层特征的细节与高层特征的语义信息通过组内融合结合，使特征在不同尺度上实现有效互补。同时，GAB还引入了深度监督生成的辅助分割信息，进一步指导特征融合，增强模型对病变区域的精准定位能力。

一 准备工作

1. 安装dynamic-network-architectures

点击链接，将其clone到本地后，进入文件夹内，pip install -e . 即可（注意-e后有个点）。

2. 生成nnUNetPlans.json文件

运行nnUNetv2_plan_and_preprocess命令，也是预处理命令，生成nnUNetPlans.json文件。

二 修改思路

1. GAB模块结构

GAB模块接受三个特征图，分别是高层次特征图、低层次特征图、mask特征图，其中，前两个只相差一层。

对高层次特征图依次进行卷积层、线性插值、沿通道维度分割 三个操作。对低层次特征图沿通道维度分割 一个操作。分割操作都均分成4份，取高层次特征图的一份、低层次特征图的一份、mask特征图，沿通道维度拼接（cat）后，通过layernorm层 + 空洞卷积层。将处理过的、新的4份特征图沿通道维度拼接（cat）后，通过layernorm层 + 卷积层。最后返回特征图。

2. 查看 nnU-Net V2 网络结构

打开nnUNet \ DATASET \ nnUNet_preprocessed \ Dataset001_ACDC \ nnUNetPlans.json文件，查看configurations --> 2d --> architecture --> network_class_name字段，默认为dynamic_network_architectures.architectures.unet.PlainConvUNet。

根据network_class_name字段，找到PlainConvUNet类所在文件：dynamic-network-architectures-main \ dynamic_network_architectures \ architectures \ unet.py

PlainConvUNet类就是nnU-Net默认的U-Net，其结构由编码器和解码器两部分组成，很标准，很常见。具体代码、结构见PlainConvUNet类

3. 替换过程

与上一篇不同，GAB模块更改了原有U-Net的跳跃连接，对特征图的传播顺序有较大改变。

由上图可知，GAB模块在下采样与上采样的连接部分，但与之前介绍的U-Net V2不同，GAB模块需要上采样（解码器）产生的mask特征图，

因此，我们将修改位置定在编码器部分。本文修改思路如下：

1. 加入GAB类
2. 在编码器部分加入GAB

三 修改网络

本次替换的一些设置

训练配置	2d
更换的网络	加入EGE-UNet的GAB模块

涉及的文件（加粗文件是要修改的文件）：

dynamic_network_architectures/
|__ architectures/
| |__ resnet.py
| |__ unet.py
| |__ vgg.py
|__ building_blocks/
| |__ helper.py
| |__ plain_conv_encoder.py
| |__ regularization.py
| |__ residual.py
| |__ residual_encoders.py
| |__ simple_conv_blocks.py
| |__ unet_decoder.py
| |__ unet_residual_decoder.py
|__ initialization/
| |__ weight_init.py

本次修改为了定位修改位置，会粘贴额外的代码用于定位，替换部分会用注释标识

1. 加入GAB类

打开dynamic-network-architectures \ dynamic_network_architectures \ building_blocks \ simple_conv_blocks.py文件，加入下面的代码：

class GAB(nn.Module):
    def __init__(self, dim_xh, dim_xl, dim_mask, k_size=3, d_list=[1,2,5,7]):
        super().__init__()
        self.pre_project = nn.Conv2d(dim_xh, dim_xl, 1)
        group_size = dim_xl // 2
        self.g0 = nn.Sequential(
            LayerNorm(normalized_shape=group_size+dim_mask, data_format='channels_first'),
            nn.Conv2d(group_size + dim_mask, group_size + dim_mask, kernel_size=3, stride=1, 
                      padding=(k_size+(k_size-1)*(d_list[0]-1))//2, 
                      dilation=d_list[0], groups=group_size + dim_mask)
        )
        self.g1 = nn.Sequential(
            LayerNorm(normalized_shape=group_size+dim_mask, data_format='channels_first'),
            nn.Conv2d(group_size + dim_mask, group_size + dim_mask, kernel_size=3, stride=1, 
                      padding=(k_size+(k_size-1)*(d_list[1]-1))//2, 
                      dilation=d_list[1], groups=group_size + dim_mask)
        )
        self.g2 = nn.Sequential(
            LayerNorm(normalized_shape=group_size+dim_mask, data_format='channels_first'),
            nn.Conv2d(group_size + dim_mask, group_size + dim_mask, kernel_size=3, stride=1, 
                      padding=(k_size+(k_size-1)*(d_list[2]-1))//2, 
                      dilation=d_list[2], groups=group_size + dim_mask)
        )
        self.g3 = nn.Sequential(
            LayerNorm(normalized_shape=group_size+dim_mask, data_format='channels_first'),
            nn.Conv2d(group_size + dim_mask, group_size + dim_mask, kernel_size=3, stride=1, 
                      padding=(k_size+(k_size-1)*(d_list[3]-1))//2, 
                      dilation=d_list[3], groups=group_size + dim_mask)
        )
        self.tail_conv = nn.Sequential(
            LayerNorm(normalized_shape=dim_xl * 2 + 4 * dim_mask, data_format='channels_first'),
            nn.Conv2d(dim_xl * 2 + 4 * dim_mask, dim_xl, 1)
        )
    def forward(self, xh, xl, mask):
        xh = self.pre_project(xh)
        xh = F.interpolate(xh, size=[xl.size(2), xl.size(3)], mode ='bilinear', align_corners=True)
        xh = torch.chunk(xh, 4, dim=1)
        xl = torch.chunk(xl, 4, dim=1)
        x0 = self.g0(torch.cat((xh[0], xl[0], mask), dim=1))
        x1 = self.g1(torch.cat((xh[1], xl[1], mask), dim=1))
        x2 = self.g2(torch.cat((xh[2], xl[2], mask), dim=1))
        x3 = self.g3(torch.cat((xh[3], xl[3], mask), dim=1))
        x = torch.cat((x0,x1,x2,x3), dim=1)
        x = self.tail_conv(x)
        return x

考虑到原版代码中的mask通道数为1，而nnU-Net V2输出的mask通道数（类别数目）不一定为1，所以对GAB类进行了适当修改，加入了dim_mask参数，用于扩展原先的1通道。

2. 在编码器部分加入GAB

打开dynamic-network-architectures \ dynamic_network_architectures \ building_blocks \ unet_decoder.py文件，修改UNetDecoder类。

首先是__init__函数，需要实例化多个GAB（一层一个），在__init__函数任意位置加入下面代码：

############################### GAB 开始
gab = []
for s in range(1, n_stages_encoder):
    gab.append(
        GAB(dim_xh=encoder.output_channels[-s], dim_xl=encoder.output_channels[-(s + 1)],
            dim_mask=self.num_classes)
    )
self.gab = nn.ModuleList(gab)
############################### GAB 结束

然后是forward函数，将该函数内的for循环内修改为如下代码：

for s in range(len(self.stages)):
    encoder_x_h, encoder_x_l = lres_input, skips[-(s+2)]
    # 首先反卷积扩大一倍
    decoder_x_l = self.transpconvs[s](encoder_x_h)
    # 再通过一个普通的卷积层获取mask
    if self.deep_supervision:
        mask_feature = self.seg_layers[s](decoder_x_l)
        seg_outputs.append(mask_feature)
    elif s == (len(self.stages) - 1):
        mask_feature = self.seg_layers[-1](decoder_x_l)
        seg_outputs.append(mask_feature)
    # 将三部分输入GAB，获取新的低层次特征图
    encoder_x_l = self.gab[s](encoder_x_h, encoder_x_l, mask_feature)
    # 逐元素相加
    decoder_x_l = torch.add(encoder_x_l, decoder_x_l)

    lres_input = decoder_x_l
    encoder_x_h = encoder_x_l

修改完毕