Attentional Feature Fusion中所提的注意力模块的代码

注意力模块的代码"/>

Attentional Feature Fusion中所提的注意力模块的代码

最近看了Attentional Feature Fusion这篇文章，对其提出的注意力机制模块很感兴趣，所以就上github找了一下，为了方便自己记录和使用，所以就复制到这里了。

源代码的github的网址：.py

论文的下载地址：.14082

MS-CAM的代码如下：

class MS_CAM(nn.Module):'''单特征 进行通道加权,作用类似SE模块'''def __init__(self, channels=64, r=4):super(MS_CAM, self).__init__()inter_channels = int(channels // r)self.local_att = nn.Sequential(nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)self.global_att = nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)self.sigmoid = nn.Sigmoid()def forward(self, x):xl = self.local_att(x)xg = self.global_att(x)xlg = xl + xgwei = self.sigmoid(xlg)return x * wei

AFF的代码如下：

class AFF(nn.Module):'''多特征融合 AFF'''def __init__(self, channels=64, r=4):super(AFF, self).__init__()inter_channels = int(channels // r)self.local_att = nn.Sequential(nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)self.global_att = nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)self.sigmoid = nn.Sigmoid()def forward(self, x, residual):xa = x + residualxl = self.local_att(xa)xg = self.global_att(xa)xlg = xl + xgwei = self.sigmoid(xlg)xo = 2 * x * wei + 2 * residual * (1 - wei)return xo# if __name__ == '__main__':
#     import os
#     os.environ['CUDA_VISIBLE_DEVICES'] = "0"
#     device = torch.device("cuda:0")
#
#     x, residual= torch.ones(8,64, 32, 32).to(device),torch.ones(8,64, 32, 32).to(device)
#     channels=x.shape[1]
#
#     model=AFF(channels=channels)
#     model=model.to(device).train()
#     output = model(x, residual)
#     print(output.shape)

### iAFF的代码如下：

class iAFF(nn.Module):'''多特征融合 iAFF'''def __init__(self, channels=64, r=4):super(iAFF, self).__init__()inter_channels = int(channels // r)# 本地注意力self.local_att = nn.Sequential(nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)# 全局注意力self.global_att = nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)# 第二次本地注意力self.local_att2 = nn.Sequential(nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)# 第二次全局注意力self.global_att2 = nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)self.sigmoid = nn.Sigmoid()def forward(self, x, residual):xa = x + residualxl = self.local_att(xa)xg = self.global_att(xa)xlg = xl + xgwei = self.sigmoid(xlg)xi = x * wei + residual * (1 - wei)xl2 = self.local_att2(xi)xg2 = self.global_att(xi)xlg2 = xl2 + xg2wei2 = self.sigmoid(xlg2)xo = x * wei2 + residual * (1 - wei2)return xo