Backbone中添加attention 之CBAM

原创

已于 2024-02-01 17:02:22 修改 · 1.4k 阅读

10 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能 #计算机视觉

于 2022-12-07 11:05:23 首次发布

本文介绍如何在ResNet网络中引入注意力机制，并提供详细的代码实现。注意力机制能够帮助模型聚焦于输入特征的重要部分，提高模型的性能。

注意力机制（Attention Mechanism）源于对人类视觉的研究。在认知科学中，由于信息处理的瓶颈，人类会选择性地关注所有信息的一部分，同时忽略其他可见的信息。上述机制通常被称为注意力机制。人类视网膜不同的部位具有不同程度的信息处理能力，即敏锐度（Acuity），只有视网膜中央凹部位具有最强的敏锐度。为了合理利用有限的视觉信息处理资源，人类需要选择视觉区域中的特定部分，然后集中关注它。例如，人们在阅读时，通常只有少量要被读取的词会被关注和处理。综上，注意力机制主要有两个方面：决定需要关注输入的哪部分；分配有限的信息处理资源给重要的部分。来至某度百科。

attention在视觉任务中有很多应用场景，每年顶会也有很多关于attention的paper。记录一下往backbone里加attention，不喜勿喷。以resnet为例，其他backbone也可以借鉴一下怎么添加attention。废话少说，直接上代码

import torch.nn as nn
import math

import torch

#3x3卷积会改变feature map 大小（当stride不等于1时），反之
def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)

#conv1x1只改变了输出的通道数
def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution """
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride,
                      bias=False)



class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        #第一个conv3x3的stride是可变的，当取2的时候，会导致特征图的size变成二分之一
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        #第二个conv3x3的stride=1,只可以改变通道数，不会改变特征图的大小。
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

        if planes == 64:
            self.globalAvgPool = nn.AvgPool2d(56, stride=1)
        elif planes == 128:
            self.globalAvgPool = nn.AvgPool2d(28, stride=1)
        elif planes == 256:
            self.globalAvgPool = nn.AvgPool2d(14, stride=1)
        elif planes == 512:
            self.globalAvgPool = nn.AvgPool2d(7, stride=1)
        self.fc1 = nn.Linear(in_features=planes, out_features=round(planes / 16))
        self.fc2 = nn.Linear(in_features=round(planes / 16), out_features=planes)
        self.sigmoid = nn.Sigmoid()

    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)

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

        original_out = out
        out = self.globalAvgPool(out)
        out = out.view(out.size(0), -1)
        out =