PyTorch 实现图像版多头注意力（Multi-Head Attention）和自注意力（Self-Attention）

本文提供一个适用于图像输入的多头注意力机制（Multi-Head Attention）PyTorch 实现，适用于 ViT、MAE 等视觉 Transformer 中的注意力计算。

模块说明

输入支持图像格式 (B, C, H, W)
内部转换为序列 (B, N, C)，其中 N = H * W
多头注意力计算：查询（Q）、键（K）、值（V）使用线性层投影
结果 reshape 回原图维度 (B, C, H, W)

多头注意力机制代码（适用于图像输入）

import torch
import torch.nn as nnclass ImageMultiHeadAttention(nn.Module):def __init__(self, embed_dim, num_heads):super(ImageMultiHeadAttention, self).__init__()assert embed_dim % num_heads == 0, "embed_dim 必须能被 num_heads 整除"self.embed_dim = embed_dimself.num_heads = num_headsself.head_dim = embed_dim // num_heads# Q, K, V 的线性映射self.q_proj = nn.Linear(embed_dim, embed_dim)self.k_proj = nn.Linear(embed_dim, embed_dim)self.v_proj = nn.Linear(embed_dim, embed_dim)# 输出映射层self.out_proj = nn.Linear(embed_dim, embed_dim)self.scale = self.head_dim ** 0.5def forward(self, x):# 输入 x: (B, C, H, W)，需要 reshape 为 (B, N, C)B, C, H, W = x.shapex = x.view(B, C, H * W).permute(0, 2, 1)  # (B, N, C)Q = self.q_proj(x)K = self.k_proj(x)V = self.v_proj(x)# 拆成多头 (B, num_heads, N, head_dim)Q = Q.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)K = K.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)V = V.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)# 注意力分数计算attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scaleattn_probs = torch.softmax(attn_scores, dim=-1)attn_out = torch.matmul(attn_probs, V)# 合并多头attn_out = attn_out.transpose(1, 2).contiguous().view(B, H * W, self.embed_dim)# 输出映射out = self.out_proj(attn_out)# 恢复回原图维度 (B, C, H, W)out = out.permute(0, 2, 1).view(B, C, H, W)return out# 测试示例
# 假设输入是一张 14x14 的特征图（类似 patch embedding 后）
img = torch.randn(4, 64, 14, 14)  # (B, C, H, W)mha = ImageMultiHeadAttention(embed_dim=64, num_heads=8)
out = mha(img)print(out.shape)  # 输出应为 (4, 64, 14, 14)

PyTorch 实现自注意力机制（Self-Attention）

本节补充自注意力机制（Self-Attention）的核心代码实现，适用于 ViT 等模型中 patch token 的注意力操作。

自注意力机制代码（Self-Attention）

import torch
import torch.nn as nnclass SelfAttention(nn.Module):def __init__(self, embed_dim):super(SelfAttention, self).__init__()self.embed_dim = embed_dimself.qkv_proj = nn.Linear(embed_dim, embed_dim * 3)self.out_proj = nn.Linear(embed_dim, embed_dim)self.scale = embed_dim ** 0.5def forward(self, x):# 输入 x: (B, N, C)B, N, C = x.shape# 一次性生成 Q, K, Vqkv = self.qkv_proj(x)  # (B, N, 3C)Q, K, V = torch.chunk(qkv, chunks=3, dim=-1)  # 各自为 (B, N, C)# 计算注意力分数attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale  # (B, N, N)attn_probs = torch.softmax(attn_scores, dim=-1)# 得到注意力加权输出attn_out = torch.matmul(attn_probs, V)  # (B, N, C)# 映射回原维度out = self.out_proj(attn_out)  # (B, N, C)return out#  测试示例
# 假设输入为 196 个 patch，每个 patch 的嵌入维度为 64
x = torch.randn(2, 196, 64)  # (B, N, C)attn = SelfAttention(embed_dim=64)
out = attn(x)print(out.shape)  # 输出应为 (2, 196, 64)

📎 拓展说明
• 本实现为单头自注意力机制
• 可用于 NLP 中的序列特征或 ViT 图像 patch 序列
• 若需改为多头注意力，只需将 embed_dim 拆成 num_heads × head_dim 并分别计算后合并

PyTorch 实现图像输入的自注意力机制（Self-Attention）

本节介绍一种适用于图像输入 (B, C, H, W) 的自注意力机制实现，适合卷积神经网络与 Transformer 的融合模块，如 Self-Attention ConvNet、BAM、CBAM、ViT 前层等。

自注意力机制（图像维度）代码

import torch
import torch.nn as nn
import torch.nn.functional as Fclass ImageSelfAttention(nn.Module):def __init__(self, in_channels):super(ImageSelfAttention, self).__init__()self.in_channels = in_channelsself.query_conv = nn.Conv2d(in_channels, in_channels // 8, kernel_size=1)self.key_conv   = nn.Conv2d(in_channels, in_channels // 8, kernel_size=1)self.value_conv = nn.Conv2d(in_channels, in_channels, kernel_size=1)self.gamma = nn.Parameter(torch.zeros(1))  # 可学习缩放因子def forward(self, x):# 输入 x: (B, C, H, W)B, C, H, W = x.size()# 生成 Q, K, Vproj_query = self.query_conv(x).view(B, -1, H * W).permute(0, 2, 1)  # (B, N, C//8)proj_key   = self.key_conv(x).view(B, -1, H * W)                      # (B, C//8, N)proj_value = self.value_conv(x).view(B, -1, H * W)                    # (B, C, N)# 注意力矩阵：Q * K^Tenergy = torch.bmm(proj_query, proj_key)         # (B, N, N)attention = F.softmax(energy, dim=-1)             # (B, N, N)# 加权求和 Vout = torch.bmm(proj_value, attention.permute(0, 2, 1))  # (B, C, N)out = out.view(B, C, H, W)# 残差连接 + 缩放因子out = self.gamma * out + xreturn out#测试用例
x = torch.randn(2, 64, 32, 32)  # 输入一张图像：B=2, C=64, H=W=32
self_attn = ImageSelfAttention(in_channels=64)
out = self_attn(x)print(out.shape)  # 输出形状应为 (2, 64, 32, 32)

• 本模块基于图像 (B, C, H, W) 进行自注意力计算
• 使用卷积进行 Q/K/V 提取，保持局部感知力
• gamma 是可学习缩放因子，用于残差连接控制注意力贡献度

自注意力中**缩放因子（scale factor）的处理，在序列维度（如 ViT）和图片维度（如 Self-Attention Conv）**中有点不一样。下面我们来详细解释一下原因，并对两种写法做一个统一和对比分析

两种缩放因子的区别

序列维度的缩放因子

scale = head_dim ** 0.5  # 或者 embed_dim ** 0.5
attn = (Q @ K.T) / scale

• 来源：Transformer 原始论文（Attention is All You Need）
• 原因：在高维向量内积中，为了避免 dot product 的结果数值过大导致梯度不稳定，需要除以 sqrt(d_k)
• 使用场景：多头注意力机制，输入是 (B, N, C)，应用在 NLP、ViT 等序列结构

图片维度（C, H, W）的注意力机制中没有缩放，或者使用 softmax 平衡

attn = softmax(Q @ K.T)   # 无 scale，或者手动调节

• 来源：Non-local Net、Self-Attention Conv、BAM 等 CNN + Attention 融合方法
• 原因：Q 和 K 都通过 1x1 conv 压缩成 C//8 或更小的维度，内积的值本身不会太大；同时图像 attention 主要用 softmax 控制权重范围
• 缩放因子的控制通常用 γ（gamma）作为残差通道缩放，不是 QK 内部的数值缩放