注意力机制——ECANet及Mobilenetv2模型应用

一、介绍
ECANet（CVPR 2020）作为一种轻量级的注意力机制，其实也是通道注意力机制的一种实现形式。其论文和开源代码为：
论文地址：https://arxiv.org/abs/1910.03151
代码：https://github.com/BangguWu/ECANet
ECA模块，去除了原来SE模块中的全连接层，直接在全局平均池化之后的特征上通过一个1D卷积进行学习。
具体的讲：通过共享相同的学习参数，通过内核大小为k的1维卷积来实现通道之间的信息交互：（一维卷积和1 × 1 卷积是不同的，一维指的是1 × k 的卷积）
ECA-Net可以插入到其他CNN网络中来增强其性能，比如：插入到ResNet、MobileNetV2中。本文主要将ECA模块加入到Mobilenetv2的残差堆叠块中。
文中同样附上SENet的嵌入代码（已注释），如有需要，可进行比较；因项目需要转换caffe模型（具体torch如何转，请看之前的博文），经测试SENet虽然转换成功，但测试时所需的caffe库不支持，所以换成ECA-Net，经转换测试，可正常出结果，且效果提升大约五个点左右。
ReLU6替换为Leakyrelu，同样是因为不支持的原因（板子太老）
二、代码
eca_module.py

import torch
from torch import nn
from torch.nn.parameter import Parameter

classeca_layer(nn.Module):"""Constructs a ECA module.
    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """def__init__(self, channel, k_size=3):super(eca_layer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1,1, kernel_size=k_size, padding=(k_size -1)//2, bias=False) 
        self.sigmoid = nn.Sigmoid()defforward(self, x):# feature descriptor on the global spatial information
        y = self.avg_pool(x)# Two different branches of ECA module
        y = self.conv(y.squeeze(-1).transpose(-1,-2)).transpose(-1,-2).unsqueeze(-1)# Multi-scale information fusion
        y = self.sigmoid(y)return x * y.expand_as(x)

eca_mobilenetv2.py

import math
import os

import torch
import torch.nn as nn
import torch.utils.model_zoo as model_zoo

from.eca_module import eca_layer

BatchNorm2d = nn.BatchNorm2d

defconv_bn(inp, oup, stride):return nn.Sequential(
        nn.Conv2d(inp, oup,3, stride,1, bias=False),
        BatchNorm2d(oup),# nn.ReLU6(inplace=True)
        nn.LeakyReLU(0.1))defconv_1x1_bn(inp, oup):return nn.Sequential(
        nn.Conv2d(inp, oup,1,1,0, bias=False),
        BatchNorm2d(oup),# nn.ReLU6(inplace=True)# nn.ReLU(inplace=True)
        nn.LeakyReLU(0.1))def_make_divisible(v, divisor, min_value=None):if min_value isNone:
        min_value = divisor
    new_v =max(min_value,int(v + divisor /2)// divisor * divisor)# Make sure that round down does not go down by more than 10%.if new_v <0.9* v:
        new_v += divisor
    return new_v

classh_sigmoid(nn.Module):def__init__(self, inplace=True):super(h_sigmoid, self).__init__()
        self.relu = nn.ReLU6(inplace=inplace)defforward(self, x):return self.relu(x +3)/6classSELayer(nn.Module):def__init__(self, channel, reduction=4):super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
                nn.Linear(channel, _make_divisible(channel // reduction,8)),
                nn.ReLU(inplace=True),
                nn.Linear(_make_divisible(channel // reduction,8), channel),
                h_sigmoid())defforward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c,1,1)return x * y

classInvertedResidual(nn.Module):def__init__(self, inp, oup, stride, expand_ratio,k_size):super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in[1,2]

        hidden_dim =round(inp * expand_ratio)
        self.use_res_connect = self.stride ==1and inp == oup

        layers =[]if expand_ratio ==1:
            layers.append(eca_layer(oup, k_size))
            self.conv = nn.Sequential(#--------------------------------------------##   进行3x3的逐层卷积，进行跨特征点的特征提取#--------------------------------------------#
                nn.Conv2d(hidden_dim, hidden_dim,3, stride,1, groups=hidden_dim, bias=False),
                BatchNorm2d(hidden_dim),# nn.ReLU6(inplace=True),
                nn.LeakyReLU(0.1),# SELayer(hidden_dim),#-----------------------------------##   利用1x1卷积进行通道数的调整#-----------------------------------#
                nn.Conv2d(hidden_dim, oup,1,1,0, bias=False),
                BatchNorm2d(oup),)else:
            layers.append(eca_layer(oup, k_size))
            self.conv = nn.Sequential(#-----------------------------------##   利用1x1卷积进行通道数的上升#-----------------------------------#
                nn.Conv2d(inp, hidden_dim,1,1,0, bias=False),
                BatchNorm2d(hidden_dim),# nn.ReLU6(inplace=True),
                nn.LeakyReLU(0.1),#--------------------------------------------##   进行3x3的逐层卷积，进行跨特征点的特征提取#--------------------------------------------#
                nn.Conv2d(hidden_dim, hidden_dim,3, stride,1, groups=hidden_dim, bias=False),
                BatchNorm2d(hidden_dim),# SELayer(hidden_dim),# nn.ReLU6(inplace=True),
                nn.LeakyReLU(0.1),#-----------------------------------##   利用1x1卷积进行通道数的下降#-----------------------------------#
                nn.Conv2d(hidden_dim, oup,1,1,0, bias=False),
                BatchNorm2d(oup),)defforward(self, x):if self.use_res_connect:return x + self.conv(x)else:return self.conv(x)classMobileNetV2(nn.Module):def__init__(self, n_class=1000, input_size=224, width_mult=1.):super(MobileNetV2, self).__init__()
        block = InvertedResidual
        input_channel =32
        last_channel =1280
        interverted_residual_setting =[# t, c, n, s[1,16,1,1],# 256, 256, 32 -> 256, 256, 16[6,24,2,2],# 256, 256, 16 -> 128, 128, 24   2[6,32,3,2],# 128, 128, 24 -> 64, 64, 32     4[6,64,4,2],# 64, 64, 32 -> 32, 32, 64       7[6,96,3,1],# 32, 32, 64 -> 32, 32, 96[6,160,3,2],# 32, 32, 96 -> 16, 16, 160     14[6,320,1,1],# 16, 16, 160 -> 16, 16, 320]assert input_size %32==0
        input_channel =int(input_channel * width_mult)
        self.last_channel =int(last_channel * width_mult)if width_mult >1.0else last_channel
        # 512, 512, 3 -> 256, 256, 32
        self.features =[conv_bn(3, input_channel,2)]for t, c, n, s in interverted_residual_setting:
            output_channel =int(c * width_mult)for i inrange(n):# 判断 ksize 值if c <96:
                    ksize =1else:
                    ksize =3#     stride = s if i == 0 else 1if i ==0:
                    self.features.append(block(input_channel, output_channel, s, expand_ratio=t, k_size=ksize))else:
                    self.features.append(block(input_channel, output_channel,1, expand_ratio=t, k_size=ksize))
                input_channel = output_channel

        self.features.append(conv_1x1_bn(input_channel, self.last_channel))
        self.features = nn.Sequential(*self.features)

        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(self.last_channel, n_class),)

        self._initialize_weights()defforward(self, x):
        x = self.features(x)

        x = x.mean(3).mean(2)
        x = self.classifier(x)return x

    def_initialize_weights(self):for m in self.modules():ifisinstance(m, nn.Conv2d):
                n = m.kernel_size[0]* m.kernel_size[1]* m.out_channels
                m.weight.data.normal_(0, math.sqrt(2./ n))if m.bias isnotNone:
                    m.bias.data.zero_()elifisinstance(m, BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()elifisinstance(m, nn.Linear):
                n = m.weight.size(1)
                m.weight.data.normal_(0,0.01)
                m.bias.data.zero_()defload_url(url, model_dir='./model_data', map_location=None):ifnot os.path.exists(model_dir):
        os.makedirs(model_dir)
    filename = url.split('/')[-1]
    cached_file = os.path.join(model_dir, filename)if os.path.exists(cached_file):return torch.load(cached_file, map_location=map_location)else:return model_zoo.load_url(url,model_dir=model_dir)defmobilenetv2(pretrained=False,**kwargs):
    model = MobileNetV2(n_class=1000,**kwargs)if pretrained:
        model.load_state_dict(load_url('https://github.com/bubbliiiing/deeplabv3-plus-pytorch/releases/download/v1.0/mobilenet_v2.pth.tar'), strict=False)return model

if __name__ =="__main__":
    model = mobilenetv2()for i, layer inenumerate(model.features):print(i, layer)