本文提供一个适用于图像输入的多头注意力机制(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)
多头注意力机制代码(适用于图像输入)
python
import torch
import torch.nn as nn
class 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_dim
self.num_heads = num_heads
self.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.5
def forward(self, x):
# 输入 x: (B, C, H, W),需要 reshape 为 (B, N, C)
B, C, H, W = x.shape
x = 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.scale
attn_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)
python
import torch
import torch.nn as nn
class SelfAttention(nn.Module):
def __init__(self, embed_dim):
super(SelfAttention, self).__init__()
self.embed_dim = embed_dim
self.qkv_proj = nn.Linear(embed_dim, embed_dim * 3)
self.out_proj = nn.Linear(embed_dim, embed_dim)
self.scale = embed_dim ** 0.5
def forward(self, x):
# 输入 x: (B, N, C)
B, N, C = x.shape
# 一次性生成 Q, K, V
qkv = 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 前层等。
自注意力机制(图像维度)代码
python
import torch
import torch.nn as nn
import torch.nn.functional as F
class ImageSelfAttention(nn.Module):
def __init__(self, in_channels):
super(ImageSelfAttention, self).__init__()
self.in_channels = in_channels
self.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, V
proj_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^T
energy = torch.bmm(proj_query, proj_key) # (B, N, N)
attention = F.softmax(energy, dim=-1) # (B, N, N)
# 加权求和 V
out = torch.bmm(proj_value, attention.permute(0, 2, 1)) # (B, C, N)
out = out.view(B, C, H, W)
# 残差连接 + 缩放因子
out = self.gamma * out + x
return 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 内部的数值缩放
💬 如果你觉得这篇整理有帮助,欢迎点赞收藏!