自注意力机制可以捕捉特征图中不同位置之间的依赖关系,它通过计算特征图中每个位置与其他所有位置的相似度,得到一个注意力权重矩阵,然后根据这个权重矩阵对特征图进行加权求和,从而实现对特征图的重新表示。在自然语言处理领域,Transformer 模型中的多头自注意力机制取得了巨大成功,后来也被广泛应用到计算机视觉领域。
理解代码可以通过观察特征图形状变化入手:自注意力机制在处理过程中会对特征图进行形状调整和矩阵运算,但最终输出的特征图形状与输入特征图形状相同。
1)输入特征图形状:假设输入特征图 x
的形状为 (batch_size, channels, height, width)
,在测试代码中为 (1, 64, 32, 32)
。
2)查询、键和值特征图的形状:[ 通过卷积 + view()]得到
q
的形状为(batch_size, height * width, in_channels // 8)
,即(1, 32 * 32, 64 // 8) = (1, 1024, 8)
。k
的形状为(batch_size, in_channels // 8, height * width)
,即(1, 8, 1024)
。v
的形状为(batch_size, in_channels, height * width)
,即(1, 64, 1024)
。
3)注意力分数和注意力权重形状:
[首先,通过矩阵乘法q*k得注意力分数attention;然后,通过softmax得到注意力权重attention]
attention = torch.bmm(q, k)
attention = self.softmax(attention)
- 计算注意力分数:
torch.bmm
是批量矩阵乘法函数,用于计算查询和键的矩阵乘法,得到注意力分数。注意力分数表示每个位置与其他所有位置之间的相关性。 - 计算注意力权重:对注意力分数应用 Softmax 函数,将其转换为注意力权重。注意力权重反映了每个位置在计算最终输出时的重要程度。
attention
的形状为(batch_size, height * width, height * width)
,即(1, 1024, 1024)
。意义:确定不同位置之间的相关性
4)注意力加权的特征图形状:
out = torch.bmm(v, attention.permute(0, 2, 1))
out = out.view(batch_size, channels, height, width)
- 计算注意力加权的值特征图:将值特征图与注意力权重的转置进行矩阵乘法,得到注意力加权的值特征图。这一步相当于根据注意力权重对值特征图进行加权求和,突出重要位置的信息。
out
在进行矩阵乘法后的形状为(batch_size, in_channels, height * width)
,即(1, 64, 1024)
。 - 恢复特征图形状:使用
view
函数将注意力加权的值特征图的形状恢复为与输入特征图相同的形状,形状恢复为(batch_size, channels, height, width)
,即(1, 64, 32, 32)
5)特征融合
out = self.gamma * out + x
return out
将注意力加权后的特征图与原始输入特征图进行融合,通过可学习的缩放因子 self.gamma
控制两者的比重。这种残差连接的方式可以避免梯度消失问题,同时保留原始输入的信息,有助于模型的训练和性能提升。
简例代码:
import torch
import torch.nn as nn
class SelfAttention(nn.Module):
def __init__(self, in_channels):
super(SelfAttention, self).__init__()
# 定义查询(Query)卷积层,将输入通道数映射到 in_channels // 8
self.query = nn.Conv2d(in_channels, in_channels // 8, kernel_size=1)
# 定义键(Key)卷积层,将输入通道数映射到 in_channels // 8
self.key = nn.Conv2d(in_channels, in_channels // 8, kernel_size=1)
# 定义值(Value)卷积层,通道数保持不变
self.value = nn.Conv2d(in_channels, in_channels, kernel_size=1)
# 可学习的缩放因子,初始化为 0
self.gamma = nn.Parameter(torch.zeros(1))
# 定义 Softmax 激活函数,用于计算注意力权重
self.softmax = nn.Softmax(dim=-1)
def forward(self, x):
# 获取输入特征图的批次大小、通道数、高度和宽度
batch_size, channels, height, width = x.size()
# 通过查询卷积层得到查询特征图,并调整形状为 (batch_size, height * width, in_channels // 8)
q = self.query(x).view(batch_size, -1, height * width).permute(0, 2, 1)
# 通过键卷积层得到键特征图,并调整形状为 (batch_size, in_channels // 8, height * width),permute 函数用于交换张量的维度。
k = self.key(x).view(batch_size, -1, height * width)
# 通过值卷积层得到值特征图,并调整形状为 (batch_size, in_channels, height * width)
v = self.value(x).view(batch_size, -1, height * width)
# 计算查询和键的矩阵乘法,得到注意力分数
attention = torch.bmm(q, k)
# 对注意力分数应用 Softmax 函数,得到注意力权重
attention = self.softmax(attention)
# 计算注意力加权的值特征图
out = torch.bmm(v, attention.permute(0, 2, 1))
# 调整形状为 (batch_size, channels, height, width)
out = out.view(batch_size, channels, height, width)
# 将注意力加权后的特征图与输入特征图进行融合
out = self.gamma * out + x
return out
# 测试代码
if __name__ == "__main__":
# 随机生成一个输入特征图,形状为 (batch_size, channels, height, width)
input_tensor = torch.randn(1, 64, 32, 32)
# 创建自注意力模块实例
self_attn = SelfAttention(in_channels=64)
# 前向传播
output = self_attn(input_tensor)
print("输入特征图形状:", input_tensor.shape)
print("输出特征图形状:", output.shape)
讨论:自注意力与空间注意力的区别
1)关注重点不同
- 自注意力机制:重点关注的是输入序列(或特征图)中不同位置之间的依赖关系,它在计算注意力权重时,会考虑所有位置之间的相互作用,以捕捉全局的长距离依赖信息,无论这些位置在空间上是否相邻。
- 空间注意力机制:主要聚焦于空间位置上的信息,它通过对特征图的空间维度进行分析,来确定不同空间位置的重要性,更侧重于挖掘空间上的局部或全局信息分布,强调的是空间位置的权重分配。
2)计算方式不同:
- 自注意力机制:通过计算查询(Query)、键(Key)和值(Value)之间的相似度来确定注意力权重。通常将输入特征图分别映射到查询、键和值空间,然后进行矩阵乘法等操作,得到注意力分数,再通过 Softmax 函数归一化得到注意力权重,计算过程涉及到较为复杂的矩阵运算和维度变换。
- 空间注意力机制:一般是通过对特征图在空间维度上进行卷积操作、池化操作或其他简单的统计运算,如计算平均值、最大值等,来获取空间位置的特征信息,进而生成空间注意力权重,计算方式相对较为直接,主要围绕空间维度进行操作
3)作用范围不同
- 自注意力机制:作用范围是整个输入序列或特征图的所有位置,它可以捕捉到非常长距离的依赖关系,即使两个位置在空间上相距很远,也能通过计算它们之间的注意力权重来建立联系,在处理具有长期依赖关系的任务,如自然语言处理中的长文本序列、图像中的全局语义关系等方面表现出色。
- 空间注意力机制:主要作用于特征图的空间维度,虽然也能关注到全局空间信息,但更侧重于对空间局部区域的信息进行加权,对于空间上相邻或相近位置的信息交互更为关注,在处理需要突出空间局部特征的任务,如目标检测中的目标定位、图像分割中的边界识别等方面有很好的效果。
4)应用场景偏好不同
- 自注意力机制:在自然语言处理领域应用广泛,如在 Transformer 架构中用于处理长序列文本,能够很好地捕捉文本中的语义依赖关系。在图像领域,对于一些需要理解图像全局语义信息、捕捉图像中不同物体之间长距离关系的任务,如场景理解、图像生成等,自注意力机制也能发挥重要作用。
- 空间注意力机制:在图像任务中应用更为普遍,特别是在图像检测、分割等任务中,能够帮助模型更好地聚焦于目标物体的空间位置,准确地识别和分割出目标物体。在一些涉及空间信息处理的计算机视觉任务,如视频动作识别、自动驾驶中的场景感知等方面也有广泛的应用。