SENet(Squeeze Excitation)、CBAM(Convolutional Block Attention Module)和CA(Coordinate Attention)

转载已于 2024-10-16 10:55:05 修改 · 974 阅读

1 ·

CC 4.0 BY-SA版权

原文链接：https://zhuanlan.zhihu.com/p/80123284

文章标签：

#深度学习

于 2022-01-02 14:38:21 首次发布

算法专栏收录该内容

475 篇文章

订阅专栏

SENet(Squeeze Excitation)

我们可以看到，已经有很多工作在空间维度上来提升网络的性能。那么很自然想到，网络是否可以从其他层面来考虑去提升性能，比如考虑特征通道之间的关系？我们的工作就是基于这一点并提出了 Squeeze-and-Excitation Networks（简称 SENet）。Squeeze 和 Excitation 是两个非常关键的操作，所以我们以此来命名。动机是希望显式地建模特征通道之间的相互依赖关系。另外，我们并不打算引入一个新的空间维度来进行特征通道间的融合，而是采用了一种全新的「特征重标定」策略。具体来说，就是通过学习的方式来自动获取到每个特征通道的重要程度，然后依照这个重要程度去提升有用的特征并抑制对当前任务用处不大的特征。

简单来说就分为三步：

Squeeze：通过Global Average pooling [n,c,h,w]->[n,c,1,1]，得到一个w矩阵
Excitation: 把这个w矩阵通过FC和Relu变成[n,c/r,1,1]，再通过FC和Sigmoid变成[n,c,1,1]
Reweight：最后用这个w矩阵来reweight输入

SENet的code如下:

code摘录自 imgclsmob/pytorch/pytorchcv/models/common.py at 68335927ba27f2356093b985bada0bc3989836b1 · osmr/imgclsmob · GitHub

import torch.nn as nn
import torch

def conv1x1(in_channels,
            out_channels,
            stride=1,
            groups=1,
            bias=False):
    return nn.Conv2d(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=1,
        stride=stride,
        groups=groups,
        bias=bias)

class SEBlock(nn.Module):

    def __init__(self,
                 channels,
                 reduction=16):
        super(SEBlock, self).__init__()
        mid_cannels = channels // reduction

        self.pool = nn.AdaptiveAvgPool2d(output_size=1)
        self.conv1 = conv1x1(
            in_channels=channels,
            out_channels=mid_cannels,
            bias=True)
        self.activ = nn.ReLU(inplace=True)
        self.conv2 = conv1x1(
            in_channels=mid_cannels,
            out_channels=channels,
            bias=True)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        w = self.pool(x)
        print('after pool shape {}'.format(w.shape))
        w = self.conv1(w)
        print('after conv1 shape {}'.format(w.shape))
        w = self.activ(w)
        print('after activ shape {}'.format(w.shape))
        w = self.conv2(w)
        print('after conv2 shape {}'.format(w.shape))
        w = self.sigmoid(w)
        print('after sigmoid shape {}'.format(w.shape))
        x = x * w
        return x

if __name__ == '__main__':
    bs,c,h,w=2,16,3,3
    x = torch.randn(bs,c,h,w)
    se = SEBlock(channels=c)
    print(se(x).shape)

'''
输出为：
after pool shape torch.Size([2, 16, 1, 1])
after conv1 shape torch.Size([2, 1, 1, 1])
after activ shape torch.Size([2, 1, 1, 1])
after conv2 shape torch.Size([2, 16, 1, 1])
after sigmoid shape torch.Size([2, 16, 1, 1])
torch.Size([2, 16, 3, 3])
'''

Convolutional Block Attention Module (CBAM)

该网络发表于2018年的CVPR。其主要思想是对特征进行空间和通道上的注意力操作，可以认为是SENet的增强版。

图中第一阶段Channel Attention Module相当于SENet中的pooling用两种方式做了两遍，代码中叫作ChannelGate；第二阶段Spatial Attention Module相当于先用ChannelPool分别用均值和最大两种pooling的方式把[n,c,h,w]压缩到[n,2,h,w]（代码中ChannelPool的部分），然后再通过一个input_channels=2,output_channels=1的卷积变成[n,1,h,w]，最后来一个sigmoid后reweight的操作。ChannelGate整体相当于SENet中pooling->mlp->relu->mlp的方式，而ChannelPool是直接做2种方式的pooling，细节可以参考下面的代码：

import torch
import math
import torch.nn as nn
import torch.nn.functional as F
 
class BasicConv(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):
        super(BasicConv, self).__init__()
        self.out_channels = out_planes
        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else None
        self.relu = nn.ReLU() if relu else None
 
    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x
 
class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)
 
class ChannelGate(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
        super(ChannelGate, self).__init__()
        self.gate_channels = gate_channels
        self.mlp = nn.Sequential(
            Flatten(),
            nn.Linear(gate_channels, gate_channels // reduction_ratio),
            nn.ReLU(),
            nn.Linear(gate_channels // reduction_ratio, gate_channels)
            )
        self.pool_types = pool_types
    def forward(self, x):
        channel_att_sum = None
        for pool_type in self.pool_types:
            if pool_type=='avg':
                avg_pool = F.avg_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( avg_pool )
            elif pool_type=='max':
                max_pool = F.max_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( max_pool )
            elif pool_type=='lp':
                lp_pool = F.lp_pool2d( x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( lp_pool )
            elif pool_type=='lse':
                # LSE pool only
                lse_pool = logsumexp_2d(x)
                channel_att_raw = self.mlp( lse_pool )
 
            if channel_att_sum is None:
                channel_att_sum = channel_att_raw
            else:
                channel_att_sum = channel_att_sum + channel_att_raw
 
        scale = F.sigmoid( channel_att_sum ).unsqueeze(2).unsqueeze(3).expand_as(x)
        return x * scale
 
def logsumexp_2d(tensor):
    tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)
    s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)
    outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()
    return outputs
 
class ChannelPool(nn.Module):
    def forward(self, x):
        # torch.max(x,1)[0] 表示在dim=1这维算最大，max返回namedtuple (values, indices) ，这里只取values
        # 整个函数是把channel维度的最大和平均给拼起来
        return torch.cat( (torch.max(x,1)[0].unsqueeze(1), torch.mean(x,1).unsqueeze(1)), dim=1 )
 
class SpatialGate(nn.Module):
    def __init__(self):
        super(SpatialGate, self).__init__()
        kernel_size = 7
        self.compress = ChannelPool()
        self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2, relu=False)
    def forward(self, x):
        x_compress = self.compress(x)
        print('after x_compress shape {}'.format(x_compress.shape))
        x_out = self.spatial(x_compress)
        print('after x_out shape {}'.format(x_out.shape))
        scale = F.sigmoid(x_out) # broadcasting
        return x * scale
 
class CBAM(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):
        super(CBAM, self).__init__()
        self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)
        self.no_spatial=no_spatial
        if not no_spatial:
            self.SpatialGate = SpatialGate()
    def forward(self, x):
        x_out = self.ChannelGate(x)
        print('after channelgate shape {}'.format(x_out.shape))
        if not self.no_spatial:
            x_out = self.SpatialGate(x_out)
        return x_out
 
if __name__ == '__main__':
    bs,c,h,w=2,16,3,3
    x = torch.randn(bs,c,h,w)
    cbam = CBAM(gate_channels=c)
    res = cbam(x)
    print(res.shape)

"""
after channelgate shape torch.Size([2, 16, 3, 3])
after x_compress shape torch.Size([2, 2, 3, 3])
after x_out shape torch.Size([2, 1, 3, 3])
torch.Size([2, 16, 3, 3])
"""

Coordinate Attention for Efficient Mobile Network Design（简称Coordinate Attention）

该网络发表于2021的CVPR，核心是拆成h,w两个维度进行reweight，只用avg pool，最终只reweight一次就可以了，参考下图：

代码摘录自CoordAttention/coordatt.py at main · houqb/CoordAttention · GitHub：

import torch
import torch.nn as nn
import math
import torch.nn.functional as F

class h_sigmoid(nn.Module):
    def __init__(self, inplace=True):
        super(h_sigmoid, self).__init__()
        self.relu = nn.ReLU6(inplace=inplace)

    def forward(self, x):
        return self.relu(x + 3) / 6

class h_swish(nn.Module):
    def __init__(self, inplace=True):
        super(h_swish, self).__init__()
        self.sigmoid = h_sigmoid(inplace=inplace)

    def forward(self, x):
        return x * self.sigmoid(x)

class CoordAtt(nn.Module):
    def __init__(self, inp, oup, reduction=32):
        super(CoordAtt, self).__init__()
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))

        mip = max(8, inp // reduction)

        self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
        self.bn1 = nn.BatchNorm2d(mip)
        self.act = h_swish()
        
        self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
        self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
        

    def forward(self, x):
        identity = x
        
        n,c,h,w = x.size()
        x_h = self.pool_h(x)
        x_w = self.pool_w(x).permute(0, 1, 3, 2)

        y = torch.cat([x_h, x_w], dim=2)
        y = self.conv1(y)
        y = self.bn1(y)
        y = self.act(y) 
        
        x_h, x_w = torch.split(y, [h, w], dim=2)
        x_w = x_w.permute(0, 1, 3, 2)

        a_h = self.conv_h(x_h).sigmoid()
        a_w = self.conv_w(x_w).sigmoid()

        out = identity * a_w * a_h

        return out