SE通道注意力机制模块

简介

论文原址:https://arxiv.org/pdf/1709.01507.pdf

在深度学习领域,提升模型的表征能力一直是一个关键的研究方向。SE(Squeeze-and-Excitation)模块是一种引入通道注意力机制的方法,旨在让神经网络更加关注对当前任务重要的特征。本文将介绍SE模块的原理、实现以及在深度学习中的应用。

SE模块结构

SE通道注意力机制模块_第1张图片

SE模块主要的两个操作就是Squeeze与Excitation操作。

首先是Squeeze操作,通过聚合跨空间维度(H × W)的特征映射来产生通道描述符,怎么理解呢?

假设有一个输入的特征映射,它的维度是H × W × C,对于每个通道,执行全局平均池化操作,具体来说,对于第i个通道,计算该通道上所有空间位置的平均值。可以通过对每个通道的所有元素取平均来实现,其中每个元素表示相应通道上的平均值,在对输入特征映射中的所有通道,都执行上述的操作,就会得到C个平均值,形成一个长度为C的向量。由此,Squeeze操作将输入的H × W × C特征映射压缩成一个C维的向量,其中每个元素表示对应通道上的平均值。这个向量可以被看作是对整个特征映射在通道维度上的"描述符"。


接着是Excitation操作,它是SE模块中的第二步,其主要目的是对上一步Squeeze得到的通道描述符进行加权,以强调重要的通道。

接着是将Squeeze得到的C维向量输入到一个全连接层(FC层),其目的是学习通道权重。这个全连接层包括一个非线性激活函数,如ReLU,以引入非线性变换。通过学习,全连接层得到的通道权重经过一个Sigmoid激活函数,将其范围限制在0到1之间。学到的权重可以被视为每个通道的激活程度。最后,将原始的C维向量与这学到的通道权重相乘,得到加权后的向量。这个加权后的向量反映了每个通道在任务中的重要性和贡献程度。这一过程使得网络能够动态调整通道的关注度,从而更有效地利用信息来完成任务。

"""
Original paper addresshttps: https://arxiv.org/pdf/1709.01507.pdf
Code originates from https://github.com/moskomule/senet.pytorch/blob/master/senet/se_module.py
"""
import torch
from torch import nn

class SELayer(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(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.expand_as(x)


if __name__ == '__main__':
    input = torch.randn(1, 64, 64, 64)
    se = SELayer(channel=64, reduction=8)
    output = se(input)
    print(output.shape)

关于这一部分的理解,我也是看了许多博主的博客,很多都是一步带过,我也是综合着才理解到原来是这样的,最后会过头来再看代码似乎有一点明悟,我的理解可以与其对的上。

这里还涉及到了一个超参数reduction,指的是 SE 模块的缩减比例。据作者论文里面说的是16是一个效果比较好的值,它表示的是通道数缩减的比例为 1/16。这个比例用于控制全连接层的通道数缩减,从而影响 SE 模块的参数量和计算复杂度,通过调整 reduction 参数,可以灵活控制通道数的缩减程度。

SE-Resnet网络

由于我也是刚学完这部分的知识,应用上面还需要进行实验,所以,还是来看看原作者提供的torch版本的网络。

import torch.nn as nn
from torchvision.models import ResNet

__all__ = ["se_resnet18", "se_resnet34", "se_resnet50", "se_resnet101", "se_resnet152"]

class SELayer(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(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.expand_as(x)

def conv3x3(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)

class SEBasicBlock(nn.Module):
    expansion = 1
    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None,
                 *, reduction=16):
        super(SEBasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes, 1)
        self.bn2 = nn.BatchNorm2d(planes)
        self.se = SELayer(planes, reduction)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.se(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class SEBottleneck(nn.Module):
    expansion = 4
    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None,
                 *, reduction=16):
        super(SEBottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.se = SELayer(planes * 4, reduction)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)
        out = self.se(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


def se_resnet18(num_classes):
    model = ResNet(SEBasicBlock, [2, 2, 2, 2], num_classes=num_classes)
    model.avgpool = nn.AdaptiveAvgPool2d(1)
    return model


def se_resnet34(num_classes):
    model = ResNet(SEBasicBlock, [3, 4, 6, 3], num_classes=num_classes)
    model.avgpool = nn.AdaptiveAvgPool2d(1)
    return model


def se_resnet50(num_classes):
    model = ResNet(SEBottleneck, [3, 4, 6, 3], num_classes=num_classes)
    model.avgpool = nn.AdaptiveAvgPool2d(1)
    return model


def se_resnet101(num_classes):
    model = ResNet(SEBottleneck, [3, 4, 23, 3], num_classes=num_classes)
    model.avgpool = nn.AdaptiveAvgPool2d(1)
    return model


def se_resnet152(num_classes):
    model = ResNet(SEBottleneck, [3, 8, 36, 3], num_classes=num_classes)
    model.avgpool = nn.AdaptiveAvgPool2d(1)
    return model

if __name__=="__main__":
    import torchsummary
    import torch
    input = torch.ones(2, 3, 224, 224).cuda()
    net = se_resnet50(num_classes=4)
    net = net.cuda()
    out = net(input)
    print(out)
    print(out.shape)
    torchsummary.summary(net, input_size=(3, 224, 224))
    # Total params: 26,031,172

这里面基本上是与torchvision中的源码相同的,只是对原本的BasicBlock与Bottleneck在最后的卷积层后加入了SE模块。

ResNet与SE-ResNet分类性能比较实验

这一部分我是在自己的数据集上进行分类的,一共四个类,每个类别大约在300张左右,ResNet添加了官方的预训练权重,SE-Resnet是重头开始训练。(本来是预计都跑50轮的,但是在跑ResNet的时候忘记了)

下面是ResNet50的训练损失记录:

SE通道注意力机制模块_第2张图片

下面是SE_ResNet50的训练损失记录:

SE通道注意力机制模块_第3张图片

这里还记录了se_resnet的错误率记录,因为是后加的功能,所以resnet50没有。

SE通道注意力机制模块_第4张图片

这里出现了过拟合的情况,不过也没关系,本身这个数据量在这里,这里我也只是想要测试它的性能效果而已。

你可能感兴趣的:(pytorch复现,SE模块,人工智能,深度学习)