【pytorch】搭建带有注意力机制的图像分类网络，实现空间注意力

一、注意力机制简介

注意力机制是深度学习中重要的技术之一，它可以使模型在处理序列数据时，更加集中地关注重要的部分，从而提高模型的性能和效率。在自然语言处理和机器翻译等任务中，注意力机制已经得到广泛应用。
在机器视觉领域，注意力机制也被广泛应用，特别是在图像分类、目标检测和图像分割等任务中。在这些任务中，注意力机制通常用来提高模型对关键区域的关注度，从而提高模型的准确性和鲁棒性。
例如，在图像分类任务中，注意力机制可以帮助模型在图像中关注到更加重要的区域，从而提高分类的准确性。在目标检测任务中，注意力机制可以帮助模型更加准确地定位目标，从而提高检测的准确性。在图像分割任务中，注意力机制可以帮助模型更加准确地分割出目标区域，从而提高分割的准确性。

二、pytorch代码实现注意力机制

我们定义了一个Attention模块，用于计算注意力权重，并将其应用于卷积层的输出。然后，我们定义了一个带有注意力机制的图像分类模型，其中包括两个卷积层和一个全连接层。在模型的前两个卷积层中，我们使用了Attention模块。最后，我们使用交叉熵损失和Adam优化器来训练模型，并在测试数据集上测试了模型的准确率。
通过这个示例代码，我们可以看到如何在PyTorch中实现注意力机制，从而在图像分类任务中提高模型的性能。

import torch
import torch.nn as nn
import torch.nn.functional as F
# 定义注意力模块
"""
首先通过一个卷积层将输入特征图x映射到一个大小为1通道数的特征图a上。
然后对映射后的特征图a进行ReLU激活函数处理，保证特征图a的数值非负。
将特征图a拉伸成一维向量，并在该向量上应用softmax函数，得到一个标准化的权重向量。
将标准化的权重向量reshape成与特征图a相同的形状，得到注意力权重矩阵a。
将原始特征图x与注意力权重矩阵a逐元素相乘，得到加权后的特征图x。
返回加权后的特征图x。
通过这种方式，Attention模块能够对输入的特征图进行加权重组合，使得模型更加关注重要的特征，从而提高模型的性能。
"""
class Attention(nn.Module):
    def __init__(self, in_channels):
        super(Attention, self).__init__()
        #输入及输出都为3通道，不改变原始图片通道数
        self.conv = nn.Conv2d(in_channels, 3, kernel_size=1, bias=False)
    def forward(self, x):
        a = F.relu(self.conv(x))
        a = F.softmax(a.view(a.size(0), -1), dim=1).view_as(a)
        x = x * a
        return x
# 构建新的网络，将resnet18的输出作为输入
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        #注意力，用于区分输入图片重要的部分
        self.attention1 = Attention(3)
        self.resnet18 = resnet18
        self.fc1 = nn.Linear(512, 256)
        self.fc2 = nn.Linear(256, 64)
        self.fc3 = nn.Linear(64, 10)
        self.fc4 = nn.Linear(10, 2)
        self.softmax = nn.Softmax(dim=1)
    def forward(self, x):
        x = self.attention1(x)
        x = self.resnet18(x)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = F.relu(self.fc4(x))
        x = self.softmax(x)
        x=x.view(-1,2)
        return x