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)

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

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)**中有点不一样。下面我们来详细解释一下原因,并对两种写法做一个统一和对比分析

两种缩放因子的区别
  1. 序列维度的缩放因子

    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 等序列结构

  1. 图片维度(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 内部的数值缩放


💬 如果你觉得这篇整理有帮助,欢迎点赞收藏!

相关推荐
0x21118 分钟前
[论文阅读]REPLUG: Retrieval-Augmented Black-Box Language Models
论文阅读·人工智能·语言模型
GeekABC41 分钟前
FastAPI系列06:FastAPI响应(Response)
开发语言·python·fastapi·web
fen_fen1 小时前
Python3:Jupyter Notebook 安装和配置
ide·python·jupyter
JOYCE_Leo161 小时前
一文详解卷积神经网络中的卷积层和池化层原理 !!
人工智能·深度学习·cnn·卷积神经网络
float_六七2 小时前
Python语言基础知识详解:分支结构控制语句
python
声声codeGrandMaster2 小时前
django之优化分页功能(利用参数共存及封装来实现)
数据库·后端·python·django
Johny_Zhao2 小时前
OpenStack 全套搭建部署指南(基于 Kolla-Ansible)
linux·python·信息安全·云计算·openstack·shell·yum源·系统运维
27669582922 小时前
海关 瑞数 后缀分析 rs
java·python·rs·瑞数·海关·瑞数后缀·后缀生成
~央千澈~2 小时前
对鸿蒙 Next 系统“成熟论”的深度剖析-优雅草卓伊凡
人工智能