Vision Transformer 图像分类实践笔记：基于 MindSpore 的实现解析

昇思MindSpore 2025-11-25 16:41:49

在计算机视觉任务中，Vision Transformer（ViT）提供了一种突破传统卷积神经网络范式的全新思路——它完全摒弃了卷积操作，转而将图像视为序列输入，并借助 Transformer 架构完成分类任务。本文基于 MindSpore 框架中的 ViT 图像分类示例，重点梳理其核心模块之一：多头注意力机制（Multi-Head Attention）的实现细节与设计思想。

多头注意力模块详解

该模块是 ViT 架构的核心组件，负责建模图像块之间的全局依赖关系。其关键在于将输入特征向量并行映射到多个“注意力头”中，每个头独立学习不同的语义子空间，最终融合各头输出以增强模型表达能力。

初始化逻辑

from mindspore import nn, ops

class Attention(nn.Cell):
    def __init__(self,
                 dim: int,
                 num_heads: int = 8,
                 keep_prob: float = 1.0,
                 attention_keep_prob: float = 1.0):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = ms.Tensor(head_dim ** -0.5)  # 缩放因子，用于稳定 softmax 输出

        # 使用一个 Dense 层一次性生成 Q、K、V
        self.qkv = nn.Dense(dim, dim * 3)
        # 注意力权重与输出的 Dropout 正则化
        self.attn_drop = nn.Dropout(p=1.0 - attention_keep_prob)
        self.out = nn.Dense(dim, dim)
        self.out_drop = nn.Dropout(p=1.0 - keep_prob)

        # 矩阵乘法操作符
        self.attn_matmul_v = ops.BatchMatMul()
        self.q_matmul_k = ops.BatchMatMul(transpose_b=True)
        self.softmax = nn.Softmax(axis=-1)

初始化阶段的关键设计包括：

维度划分：将输入维度 dim均匀分配给 num_heads个注意力头，每头处理 head_dim = dim // num_heads维度的特征。
缩放因子：引入 $\frac{1}{\sqrt{d_k}}$ 对点积结果进行缩放，避免高维空间中内积值过大导致 softmax 梯度饱和。
QKV 合并生成：通过单个全连接层一次性输出三倍维度的张量，再拆分为 Query、Key 和 Value，兼顾简洁性与计算效率。
双重 Dropout：分别对注意力权重（attn_drop）和最终输出（out_drop）施加随机失活，提升泛化能力。

前向传播流程

def construct(self, x):
    b, n, c = x.shape  # batch_size, sequence_length, embedding_dim

    # 生成 Q、K、V 并重塑为多头格式
    qkv = self.qkv(x)
    qkv = ops.reshape(qkv, (b, n, 3, self.num_heads, c // self.num_heads))
    qkv = ops.transpose(qkv, (2, 0, 3, 1, 4))  # (3, B, H, N, D)
    q, k, v = ops.unstack(qkv, axis=0)         # 拆分为三个独立张量

    # 计算注意力分数
    attn = self.q_matmul_k(q, k)               # Q @ K^T
    attn = ops.mul(attn, self.scale)           # 缩放
    attn = self.softmax(attn)                  # 归一化为权重
    attn = self.attn_drop(attn)                # 随机丢弃部分注意力连接

    # 加权聚合 Value
    out = self.attn_matmul_v(attn, v)          # (B, H, N, D)

    # 重组多头输出
    out = ops.transpose(out, (0, 2, 1, 3))     # (B, N, H, D)
    out = ops.reshape(out, (b, n, c))          # 合并所有头：(B, N, C)
    out = self.out(out)                        # 线性投影
    out = self.out_drop(out)                   # 输出 Dropout

    return out

前向过程可分为三个阶段：

QKV 重构与多头分离
输入张量经线性变换后被重塑为五维结构 (B, N, 3, H, D)，再通过转置与拆分，得到形状为 (B, H, N, D)的 Q、K、V。这种布局天然支持批量并行计算，且保证各注意力头互不干扰。
缩放点积注意力计算
利用 BatchMatMul高效完成矩阵乘法：先计算 $QK^\top$ 得到原始相似度，再乘以缩放因子 $\frac{1}{\sqrt{d_k}}$，随后经 Softmax 转换为概率分布。此步骤确保了数值稳定性，并通过 Dropout 引入随机性以防止过拟合。
输出融合与投影
各头输出拼接回原始维度后，通过一个可学习的线性层进行特征整合，最后再次应用 Dropout。这种“分而治之、再合而用之”的策略，使模型能同时捕捉局部与全局、低层与高层的多样化语义信息。

Vision Transformer 中的注意力机制不仅体现了 Transformer 架构的通用性，也展示了其在视觉任务中的强大适应能力。通过 MindSpore 的高效算子组合，该实现兼顾了代码简洁性与运行性能，为理解 ViT 的底层运作提供了清晰的技术路径。

...全文