模型拆解（二）：GeleNet

---

一、GeleNet

  论文：Salient Object Detection in Optical Remote Sensing Images Driven by Transformer（基于Transformer的光学遥感图像中的显著目标检测）
  论文链接：Salient Object Detection in Optical Remote Sensing Images Driven by Transformer
  论文代码：Github
  博客链接：Salient Object Detection in Optical Remote Sensing Images Driven by Transformer

1.1编码器：PVT-v2-b2

  使用由四个$Transformer$编码器构成的PVT作为骨干，输入图像大小为$3\times 352\times 352$，可生成四个基本全局特征${\hat{f}}_t^i$，${\hat{f}}_t^i\in R^{c_i\times h_i\times w_i}$，其中， c i ∈ { 64 , 128 , 320 , 512 } , h i / w i = 352 2 i + 1 c_i∈\{64,128,320,512\},h_i/w_i=\frac{352}{2^{i+1}} ci​∈{64,128,320,512},hi​/wi​=2i+1352​，用于提取具有全局长距离依赖性的四级基本特征嵌入。网络架构：

• class GeleNet(nn.Module)中的相关代码：

class GeleNet(nn.Module):def __init__(self, channel=32):super(GeleNet, self).__init__()#定义编码器结构,并加载预训练的PVTv2模型self.backbone = pvt_v2_b2()  # [64, 128, 320, 512]path = './model/pvt_v2_b2.pth'save_model = torch.load(path)model_dict = self.backbone.state_dict()state_dict = {k: v for k, v in save_model.items() if k in model_dict.keys()}model_dict.update(state_dict)self.backbone.load_state_dict(model_dict)...def forward(self, x):#获取编码器输出的四张特征图pvt = self.backbone(x)x1 = pvt[0] # 64x88x88x2 = pvt[1] # 128x44x44x3 = pvt[2] # 320x22x22x4 = pvt[3] # 512x11x11

• class pvt_v2_b2()中通过在初始化函数__init__()中调用调用PVT-v2架构。

@register_model
class pvt_v2_b2(PyramidVisionTransformerImpr):def __init__(self, **kwargs):super(pvt_v2_b2, self).__init__(patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1],drop_rate=0.0, drop_path_rate=0.1)

• class PyramidVisionTransformerImpr(nn.Module)中实现了对应结构。

class PyramidVisionTransformerImpr(nn.Module):def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dims=[64, 128, 256, 512],num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0.,attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]):super().__init__()self.num_classes = num_classesself.depths = depths#四个重叠嵌入层self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_chans=in_chans,embed_dim=embed_dims[0])self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_chans=embed_dims[0],embed_dim=embed_dims[1])self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_chans=embed_dims[1],embed_dim=embed_dims[2])self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_chans=embed_dims[2],embed_dim=embed_dims[3])#定义随机深度衰减规则dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule#定义第一个Transformer Encodercur = 0self.block1 = nn.ModuleList([Block(dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale,drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,sr_ratio=sr_ratios[0])for i in range(depths[0])])self.norm1 = norm_layer(embed_dims[0])#定义第二个Transformer Encodercur += depths[0]self.block2 = nn.ModuleList([Block(dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale,drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,sr_ratio=sr_ratios[1])for i in range(depths[1])])self.norm2 = norm_layer(embed_dims[1])#定义第三个Transformer Encodercur += depths[1]self.block3 = nn.ModuleList([Block(dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale,drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,sr_ratio=sr_ratios[2])for i in range(depths[2])])self.norm3 = norm_layer(embed_dims[2])#定义第四个Transformer Encodercur += depths[2]self.block4 = nn.ModuleList([Block(dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale,drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,sr_ratio=sr_ratios[3])for i in range(depths[3])])self.norm4 = norm_layer(embed_dims[3])# classification head# self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity()#初始化参数self.apply(self._init_weights)#初始化线性层、归一化层、卷积层参数def _init_weights(self, m):if isinstance(m, nn.Linear):trunc_normal_(m.weight, std=.02)if isinstance(m, nn.Linear) and m.bias is not None:nn.init.constant_(m.bias, 0)elif isinstance(m, nn.LayerNorm):nn.init.constant_(m.bias, 0)nn.init.constant_(m.weight, 1.0)elif isinstance(m, nn.Conv2d):fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channelsfan_out //= m.groupsm.weight.data.normal_(0, math.sqrt(2.0 / fan_out))if m.bias is not None:m.bias.data.zero_()def init_weights(self, pretrained=None):if isinstance(pretrained, str):logger = 1#load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)#重置每个块的丢弃路径概率def reset_drop_path(self, drop_path_rate):dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))]cur = 0for i in range(self.depths[0]):self.block1[i].drop_path.drop_prob = dpr[cur + i]cur += self.depths[0]for i in range(self.depths[1]):self.block2[i].drop_path.drop_prob = dpr[cur + i]cur += self.depths[1]for i in range(self.depths[2]):self.block3[i].drop_path.drop_prob = dpr[cur + i]cur += self.depths[2]for i in range(self.depths[3]):self.block4[i].drop_path.drop_prob = dpr[cur + i]#冻结第一个补丁嵌入层的参数def freeze_patch_emb(self):self.patch_embed1.requires_grad = False#定义不需要权重衰减的参数集合(所有位置嵌入+class token)@torch.jit.ignoredef no_weight_decay(self):return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'}  # has pos_embed may be better#返回分类器头def get_classifier(self):return self.head#重置分类头以适应新的类别数def reset_classifier(self, num_classes, global_pool=''):self.num_classes = num_classesself.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()# def _get_pos_embed(self, pos_embed, patch_embed, H, W):#     if H * W == self.patch_embed1.num_patches:#         return pos_embed#     else:#         return F.interpolate(#             pos_embed.reshape(1, patch_embed.H, patch_embed.W, -1).permute(0, 3, 1, 2),#             size=(H, W), mode="bilinear").reshape(1, -1, H * W).permute(0, 2, 1)#定义前向传播特征提取的函数def forward_features(self, x):B = x.shape[0]outs = []# stage 1#使用第一个补丁嵌入层处理输入x,获取特征图x及对应的高H和宽Wx, H, W = self.patch_embed1(x)for i, blk in enumerate(self.block1):x = blk(x, H, W)#归一化x = self.norm1(x)#调整为(B,embed_dim,H,W)x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()outs.append(x)# stage 2x, H, W = self.patch_embed2(x)for i, blk in enumerate(self.block2):x = blk(x, H, W)x = self.norm2(x)x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()outs.append(x)# stage 3x, H, W = self.patch_embed3(x)for i, blk in enumerate(self.block3):x = blk(x, H, W)x = self.norm3(x)x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()outs.append(x)# stage 4x, H, W = self.patch_embed4(x)for i, blk in enumerate(self.block4):x = blk(x, H, W)x = self.norm4(x)x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()outs.append(x)#outs列表包含四个阶段的输出特征图return outsdef forward(self, x):x = self.forward_features(x)return x

1.3D-SWSAM：方向-置换加权空间注意力模块

  光学遥感图像中的显著对象往往有一定的方向性，传统卷积并不能很好提取这一特征，对此提出方向-置换加权空间注意力模块（$D-SWSAM$）。首先对不同方向进行定向卷积，再通过通道置换将方向信息均匀融合到每个子特征中，再针对每个子特征生成对应的局部空间注意力图，然后采用加权融合操作生成最终的空间注意力图，以实现一致增强。
  在class GeleNet(nn.Module)中通过两个模块实现：

class GeleNet(nn.Module):def __init__(self, channel=32):super(GeleNet, self).__init__()...#方向卷积self.dirConv = DirectionalConvUnit(channel)#置换加权空间注意力模块SWSAMself.DSWSAM_1 = SWSAM(channel)...

• class DirectionalConvUnit(nn.Module)中实现了方向卷积。

class DirectionalConvUnit(nn.Module):def __init__(self, channel):super(DirectionalConvUnit, self).__init__()self.h_conv = nn.Conv2d(channel, channel // 4, (1, 5), padding=(0, 2))self.w_conv = nn.Conv2d(channel, channel // 4, (5, 1), padding=(2, 0))# leading diagonalself.dia19_conv = nn.Conv2d(channel, channel // 4, (5, 1), padding=(2, 0))# reverse diagonalself.dia37_conv = nn.Conv2d(channel, channel // 4, (1, 5), padding=(0, 2))def forward(self, x):#依次进行四个方向的卷积操作x1 = self.h_conv(x)x2 = self.w_conv(x)x3 = self.inv_h_transform(self.dia19_conv(self.h_transform(x)))x4 = self.inv_v_transform(self.dia37_conv(self.v_transform(x)))#将结果concat并返回x = torch.cat((x1, x2, x3, x4), 1)return x# Code from "CoANet- Connectivity Attention Network for Road Extraction From Satellite Imagery", and we modified the codedef h_transform(self, x):shape = x.size()x = torch.nn.functional.pad(x, (0, shape[-2]))x = x.reshape(shape[0], shape[1], -1)[..., :-shape[-2]]x = x.reshape(shape[0], shape[1], shape[2], shape[2]+shape[3]-1)return xdef inv_h_transform(self, x):shape = x.size()x = x.reshape(shape[0], shape[1], -1).contiguous()x = torch.nn.functional.pad(x, (0, shape[-2]))x = x.reshape(shape[0], shape[1], shape[2], shape[3]+1)x = x[..., 0: shape[3]-shape[2]+1]return xdef v_transform(self, x):x = x.permute(0, 1, 3, 2)shape = x.size()x = torch.nn.functional.pad(x, (0, shape[-2]))x = x.reshape(shape[0], shape[1], -1)[..., :-shape[-2]]x = x.reshape(shape[0], shape[1], shape[2], shape[2]+shape[3]-1)return x.permute(0, 1, 3, 2)def inv_v_transform(self, x):x = x.permute(0, 1, 3, 2)shape = x.size()x = x.reshape(shape[0], shape[1], -1).contiguous()x = torch.nn.functional.pad(x, (0, shape[-2]))x = x.reshape(shape[0], shape[1], shape[2], shape[3]+1)x = x[..., 0: shape[3]-shape[2]+1]return x.permute(0, 1, 3, 2)

• class SWSAM(nn.Module)中实现了置换加权空间注意力模块。传统的空间注意力机制CBAM对所有通道进行全局最大池化和全局平均池化，以全局方式生成空间注意力图谱，这可能会产生不充分的空间注意力图。而分组注意力机制SGE将特征分割成若干子集，并根据每个子特征生成特定的空间注意力图，以进行单独增强。但其只考虑每个子特征的注意力，忽略了不同子特征之间注意力的一致性，导致分组增强的特征缺乏一致性，并不适合SOD任务。而$SWSAM$模块对每个子特征生成局部空间注意力图，然后采用加权融合操作生成最终的空间注意力图，以实现一致增强。

# SWSAM: Shuffle Weighted Spatial Attention Module
class SWSAM(nn.Module):def __init__(self, channel=32): # group=8, branch=4, group x branch = channelsuper(SWSAM, self).__init__()self.SA1 = SpatialAttention()self.SA2 = SpatialAttention()self.SA3 = SpatialAttention()self.SA4 = SpatialAttention()self.weight = nn.Parameter(torch.ones(4, dtype=torch.float32), requires_grad=True)self.sa_fusion = nn.Sequential(BasicConv2d(1, 1, 3, padding=1),nn.Sigmoid())def forward(self, x):#通道置换x = channel_shuffle(x, 4)#特征分割x1, x2, x3, x4 = torch.split(x, 8, dim = 1)#依次生成空间注意力图s1 = self.SA1(x1)s2 = self.SA1(x2)s3 = self.SA1(x3)s4 = self.SA1(x4)#生成相应可训练的权重参数nor_weights = F.softmax(self.weight, dim=0)#融合空间注意力图s_all = s1 * nor_weights[0] + s2 * nor_weights[1] + s3 * nor_weights[2] + s4 * nor_weights[3]#用空间注意力图加强原始特征,并使用残差连接x_out = self.sa_fusion(s_all) * x + xreturn x_out

1.4KTM：知识转移模块

  SOD任务中，两个特征的乘积可以揭示两个特征共存的重要信息，有利于协同识别对象。两个特征的求和可以全面地捕捉两个特征所包含的信息。KTM模块采用注意力机制模拟特征的乘积、求和操作来识别和阐述特征图中的突出对象。代码实现：

class KTM(nn.Module):def __init__(self, channel=32):super(KTM, self).__init__()self.query_conv = nn.Conv2d(channel, channel // 2, kernel_size=1)self.key_conv = nn.Conv2d(channel, channel // 2, kernel_size=1)self.value_conv_2 = nn.Conv2d(channel, channel, kernel_size=1)self.value_conv_3 = nn.Conv2d(channel, channel, kernel_size=1)self.gamma_2 = nn.Parameter(torch.zeros(1))self.gamma_3 = nn.Parameter(torch.zeros(1))self.softmax = Softmax(dim=-1)# following DANetself.conv_2 = nn.Sequential(BasicConv2d(channel, channel, 3, padding=1),nn.ReLU(),nn.Dropout2d(0.1, False),nn.Conv2d(channel, channel, 1))self.conv_3 = nn.Sequential(BasicConv2d(channel, channel, 3, padding=1),nn.ReLU(),nn.Dropout2d(0.1, False),nn.Conv2d(channel, channel, 1))self.conv_out = nn.Sequential(nn.Dropout2d(0.1, False),nn.Conv2d(channel, channel, 1))def forward(self, x2, x3): # V#f_sumx_sum = x2 + x3 # Q#f_prox_mul = x2 * x3 # K"""inputs :x : input feature maps( B X C X H X W)returns :out : attention value + input featureattention: B X (HxW) X (HxW)"""m_batchsize, C, height, width = x_sum.size()proj_query = self.query_conv(x_sum).view(m_batchsize, -1, width * height).permute(0, 2, 1)proj_key = self.key_conv(x_mul).view(m_batchsize, -1, width * height)energy = torch.bmm(proj_query, proj_key)attention = self.softmax(energy)proj_value_2 = self.value_conv_2(x2).view(m_batchsize, -1, width * height)proj_value_3 = self.value_conv_3(x3).view(m_batchsize, -1, width * height)out_2 = torch.bmm(proj_value_2, attention.permute(0, 2, 1))out_2 = out_2.view(m_batchsize, C, height, width)out_2 = self.conv_2(self.gamma_2 * out_2 + x2)out_3 = torch.bmm(proj_value_3, attention.permute(0, 2, 1))out_3 = out_3.view(m_batchsize, C, height, width)out_3 = self.conv_3(self.gamma_3 * out_3 + x3)x_out = self.conv_out(out_2 + out_3)return x_out

1.5解码器模块

  受级联部分解码器的启发，设计了新型级联部分解码器作为显著性预测器生成预测图。

class PDecoder(nn.Module):def __init__(self, channel):super(PDecoder, self).__init__()self.relu = nn.ReLU(True)self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)self.conv_upsample1 = BasicConv2d(channel, channel, 3, padding=1)self.conv_upsample2 = BasicConv2d(channel, channel, 3, padding=1)self.conv_upsample3 = BasicConv2d(channel, channel, 3, padding=1)self.conv_upsample4 = BasicConv2d(channel, channel, 3, padding=1)self.conv_upsample5 = BasicConv2d(2*channel, 2*channel, 3, padding=1)self.conv_concat2 = BasicConv2d(2*channel, 2*channel, 3, padding=1)self.conv_concat3 = BasicConv2d(3*channel, 3*channel, 3, padding=1)self.conv4 = BasicConv2d(3*channel, 3*channel, 3, padding=1)self.conv5 = nn.Conv2d(3*channel, 1, 1)def forward(self, x1, x2, x3): # x1: 32x11x11, x2: 32x22x22, x3: 32x88x88,#接受来自D-SWSAM、KTM、SWSAM的输入x1_1 = x1 # 32x11x11x2_1 = self.conv_upsample1(self.upsample(x1)) * x2 # 32x22x22x3_1 = self.conv_upsample2(self.upsample(self.upsample(self.upsample(x1)))) * self.conv_upsample3(self.upsample(self.upsample(x2))) * x3 # 32x88x88x2_2 = torch.cat((x2_1, self.conv_upsample4(self.upsample(x1_1))), 1) # 32x22x22x2_2 = self.conv_concat2(x2_2)x3_2 = torch.cat((x3_1, self.conv_upsample5(self.upsample(self.upsample(x2_2)))), 1) # 32x88x88x3_2 = self.conv_concat3(x3_2)x = self.conv4(x3_2)x = self.conv5(x) # 1x88x88return x