注意力模块的代码"/>
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
更多推荐
Attentional Feature Fusion中所提的注意力模块的代码
发布评论