cv中多头自注意力机制与多头注意力机制

### 多头自注意力机制在计算机视觉中的代码实现 多头自注意力机制已经在多个视觉任务中展现出卓越的效果，尤其是在 Vision Transformer 中的应用。下面是一个简单的 Python 实现例子，展示了如何构建一个多头自注意力层并将其应用于二维图像数据。 ```python import torch import torch.nn as nn import math class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadSelfAttention, self).__init__() assert d_model % num_heads == 0, "d_model must be divisible by num_heads" self.d_k = d_model // num_heads self.num_heads = num_heads self.query = nn.Linear(d_model, d_model) self.key = nn.Linear(d_model, d_model) self.value = nn.Linear(d_model, d_model) self.fc_out = nn.Linear(d_model, d_model) def forward(self, x): N, C, H, W = x.shape x = x.view(N, C, -1).permute(0, 2, 1) # Reshape from [N,C,H,W] -> [N,H*W,C] Q = self.query(x).view(N, -1, self.num_heads, self.d_k).transpose(1, 2) K = self.key(x).view(N, -1, self.num_heads, self.d_k).transpose(1, 2) V = self.value(x).view(N, -1, self.num_heads, self.d_k).transpose(1, 2) energy = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) attention = torch.softmax(energy, dim=-1) out = torch.matmul(attention, V).transpose(1, 2).contiguous().view(N, -1, C) out = self.fc_out(out.permute(0, 2, 1)).view(N, C, H, W) return out # Example usage: if __name__ == "__main__": batch_size = 4 channels = 512 height = width = 32 sample_input = torch.randn(batch_size, channels, height, width) mhsa_layer = MultiHeadSelfAttention(channels, num_heads=8) output = mhsa_layer(sample_input) print(f'Output shape: {output.shape}') ``` 在这个实现中，输入张量 `x` 被视为一系列特征向量组成的序列，其中每个位置对应于原始图像的一个空间位置[^1]。通过线性变换生成查询（Q）、键（K）和值（V），接着计算这些矩阵之间的相似度得分来形成注意力权重分布，并最终加权求和得到输出表示。此过程重复多次以构成不同的“头部”，从而允许模型关注不同部分的信息[^3]。

### 香草多头自注意力机制的概念与实现 香草多头自注意力机制（Vanilla Multi-Head Self-Attention Mechanism）是一种广泛应用于深度学习中的技术，尤其在自然语言处理领域取得了显著的成功。它通过并行计算多个注意力头来捕捉输入序列的不同子空间特征，从而增强模型的表现力。 #### 多头自注意力机制的核心概念 多头自注意力机制基于单头注意力扩展而来，其核心思想是将输入向量映射到不同的线性变换空间中，并分别计算这些空间内的注意力权重[^1]。具体来说，对于给定的查询（Query）、键（Key）和值（Value），每个多头会独立地执行以下操作： 1. **线性投影**：每个头部会对 Query、Key 和 Value 进行一次线性变换，得到对应的 Q'、K' 和 V'。 2. **缩放点积注意力**：利用缩放点积公式 \( \text{Attention}(Q', K', V') = \text{softmax}(\frac{Q'(K'^T)}{\sqrt{d_k}})V' \)，其中 \( d_k \) 是 Key 的维度大小。 3. **拼接与最终投影**：各头部的结果会被拼接在一起并通过另一个线性层进行转换，形成最终输出。 这种设计允许网络关注不同位置的信息组合方式，在诸如 Transformer 架构中有广泛应用案例。 #### 实现代码示例 以下是使用 PyTorch 编写的简单版本 Vanilla Multi-Head Self-Attention Module 的实现： ```python import torch import torch.nn as nn import math class MultiHeadSelfAttention(nn.Module): def __init__(self, embed_dim, num_heads, dropout=0.1): super(MultiHeadSelfAttention, self).__init__() assert embed_dim % num_heads == 0 self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim) self.o_proj = nn.Linear(embed_dim, embed_dim) self.dropout_layer = nn.Dropout(dropout) def forward(self, x): batch_size, seq_len, _ = x.size() qkv = self.qkv_proj(x).reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) query, key, value = qkv.unbind(dim=0) attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(self.head_dim) attention_probs = torch.softmax(attention_scores, dim=-1) attention_probs = self.dropout_layer(attention_probs) context_vector = (attention_probs @ value).transpose(1, 2).contiguous().view(batch_size, seq_len, -1) output = self.o_proj(context_vector) return output ``` 上述代码定义了一个标准的多头自注意模块，支持任意嵌入维度以及指定数量的头数配置。 #### 应用场景拓展 除了 NLP 中的经典应用外，该机制也被引入计算机视觉任务里，比如用于图像分类或者单目深度估计等问题解决上[^2]。某些研究还提出了时间解耦与时空分离型自我专注组件结合常规前馈神经网路构成两种独特变压器结构形式[^3]。