Coord.Attention 协调注意力机制

此为我对于协调注意力机制的一些理解，paper来源于2021年CVPR 

import torch
import torch.nn as nn
#---------------------------------------------------#
#CA模块这个类的定义
#其参数共有三个分别为特征图像的高、宽以及通道数
#在经过CA模块的前后，特征图像的通道数并不会发生变化
#使用池化不会造成数据矩阵深度的改变，只会在高度和宽带上降低，达到降维的目的
#---------------------------------------------------#
class CA_Block(nn.Module):
    def __init__(self, channels, h, w, reduction=16):
        super(CA_Block, self).__init__()

        self.h = h
        self.w = w
        self.avg_pool_x = nn.AdaptiveAvgPool2d((h, 1))    #在x轴进行平均池化操作，x轴即为水平方向w，进而使w的值变为1
        self.avg_pool_y = nn.AdaptiveAvgPool2d((1, w))    #在y轴进行平均池化操作，y轴为垂直方向h，进而使h的值变为1
        self.conv_1x1 = nn.Conv2d(in_channels=channels, out_channels=channels // reduction, kernel_size=1, stride=1,
                                  bias=False)                         #图中的r即为reduction，进而使其输出的特征图像的通道数变为原先的1/16

        self.relu = nn.ReLU()    #relu激活函数
        self.bn = nn.BatchNorm2d(channels // reduction)   #二维的正则化操作

        self.F_h = nn.Conv2d(in_channels=channels // reduction, out_channels=channels, kernel_size=1, stride=1,
                             bias=False)              #将垂直方向上的通道数通过卷积来将其复原
        self.F_w = nn.Conv2d(in_channels=channels // reduction, out_channels=channels, kernel_size=1, stride=1,
                             bias=False)              #将水平方向上的通道数通过卷积来将其复原

        self.sigmoid_h = nn.Sigmoid()          #定义的sigmoid方法
        self.sigmoid_w = nn.Sigmoid()

    def forward(self, x):
        x_h = self.avg_pool_x(x).permute(0, 1, 3, 2) # 16，1024，13，1 ->16,1024,1,13
        x_w = self.avg_pool_y(x)  #16，1024，1，13
        #现在x_h以及x_w的shape均为16，1024，1，13    16，1024，1，13->16,1024,1,26 经过conv之后，成为16，64，1，26 过relu提升非线性因素，即为non-linear
        x_cat_conv_relu = self.relu(self.conv_1x1(torch.cat((x_h, x_w), 3))) #其shape为16，64，1，26   #在维度3上对两个tensor进行一个cat操作，之后过1*1的卷积，再过relu
                                                                             #具体的即为将维度3上的13与13相加，同时通过卷积调整其通道数为64，进而过relu提升其非线性因素，即为non-linear
        x_cat_conv_split_h, x_cat_conv_split_w = x_cat_conv_relu.split([self.h, self.w], 3)   #按照维度3以及h和w的值将这个张量分开
        #print(x_cat_conv_split_h.shape)    #对于垂直方向，其输出的shape仍为16，64，1，13
        #print(x_cat_conv_split_w.shape)    #对于水平方向，其输出的shape仍为16，64，1，13
        s_h = self.sigmoid_h(self.F_h(x_cat_conv_split_h.permute(0, 1, 3, 2)))  #对于垂直方向的，为先进行一个转置，之后通过卷积达到所原先的通道数，再过sigmoid进行归一化处理
        #print(s_h.shape)  #为16，1024，13，1
        s_w = self.sigmoid_w(self.F_w(x_cat_conv_split_w)) #对于水平方向，使用卷积来达到原先的通道数，之后进行归一化的处理
        #print(s_w.shape)  #为16，1024，1，13
        out = x * s_h.expand_as(x) * s_w.expand_as(x) #作用为进行张量形式的复原 ，将(16，1024，13，1)与(16，1024，1，13)相乘，可得到16，1024，13，13
        return out

F=torch.randn(16,1024,13,13)
print('As Begin!!')
print(F.shape)
CA=CA_Block(1024,13,13)
F=CA(F)
print('After Change!!')
print(F.shape)