1. 环境准备与任务总览嘿朋友们今天咱们来啃一块硬骨头——CS231n作业三里的Transformer图像描述模型。我知道一看到“Transformer”、“图像描述”这些词很多刚入门的朋友可能就有点发怵感觉这玩意儿太复杂代码量肯定巨大。别担心我当年做这个作业的时候也是这么想的但跟着一步步拆解下来你会发现它其实逻辑非常清晰而且PyTorch实现起来并没有想象中那么可怕。这个作业的目标很明确用Transformer模型让AI学会看图说话。具体来说就是给你一张图片比如一只猫在沙发上睡觉模型要能生成一句描述性的文字“A cat is sleeping on a sofa.” 这属于典型的视觉-语言多模态任务。作业的核心是让你从零开始亲手搭建Transformer的几个关键模块然后把它们像搭积木一样组合起来最终训练出一个能工作的模型。这比直接调用torch.nn.Transformer要有趣得多也深刻得多因为你能看清每一块“积木”是怎么工作的。我们先来看看需要准备些什么。你需要一个Python环境我强烈建议使用Conda来管理这样可以避免各种包版本冲突的“玄学”问题。作业推荐使用PyTorch 2.6.0版本这里有个小坑不同版本的PyTorchtorch.manual_seed(231)设置的随机种子产生的随机数序列可能不同。为了保证你代码的运行结果能和作业自带的数值检查numerical check对上最好严格按照要求来安装指定版本。数据集方面作业使用的是经典的COCO Captioning数据集里面包含了大量图片和对应的人工标注描述。你需要提前下载并解压到指定目录。这些准备工作看似琐碎但却是保证后续实验能顺利复现的关键我在这上面可没少踩坑。2. 理解Transformer的核心注意力机制在动手写代码之前我们必须把Transformer的“心脏”——注意力机制——给吃透。很多人一上来就被“Query, Key, Value”和“Scaled Dot-Product”这些术语绕晕了。咱们换个方式理解想象一下你在看一幅画比如《蒙娜丽莎》。Query查询就是你当前正在思考的问题比如“这幅画的主角是谁”。Key键就是画布上各个区域的特征标签比如“脸部区域”、“背景山水”、“手的姿势”。Value值就是这些区域所包含的具体信息内容。注意力机制的工作流程就像你的大脑在处理这个问题计算相关性你的问题Query会和画布上每个区域的标签Key进行比对计算出一个“相关度分数”。显然“脸部区域”这个Key和“主角是谁”这个Query的相关度会非常高。归一化权重通过Softmax函数把这些相关度分数转换成一组权重总和为1。相关度高的区域如脸部会获得接近1的权重不相关的区域如背景角落权重接近0。加权求和最后用这组权重对各个区域的具体信息Value进行加权求和。最终你得到的“输出”就是高度集中于“脸部区域信息”的合成信息从而得出“主角是蒙娜丽莎”这个结论。这就是自注意力Self-Attention的精髓让序列中的每个元素比如一句话里的每个词或一张图分成的每个图像块都能“看到”序列中的所有其他元素并根据相关性动态地聚合信息。而缩放点积注意力Scaled Dot-Product Attention就是实现上述过程的数学公式Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。除以sqrt(d_k)这个缩放因子非常重要是为了防止当特征维度d_k较大时点积结果过大导致Softmax梯度消失。2.1 为什么需要多头注意力Multi-Head Attention只用一个“大脑”单头来看画可能只能关注到“主角是谁”。但如果我们有多个并行的“大脑”多头呢一个头可能专门看颜色和光影发现画面整体偏暗另一个头可能分析构图和线条注意到人物的三角形构图第三个头可能聚焦于细节纹理比如纱巾的质感。每个头都从同一个输入出发但通过不同的参数矩阵不同的W_Q, W_K, W_V将其投影到不同的“子空间”从而关注不同类型的信息。在代码里我们不是训练多个独立的模型而是将模型的隐藏维度d_model平均分成num_heads份每一份作为一个“头”在各自的低维空间里并行计算注意力。最后把所有头的输出拼接起来再通过一个线性层W_O融合。这样做的好处是极大地增强了模型的表征能力让它能同时捕捉词语之间多种多样的关系比如语法结构、语义关联、指代关系等。我在实际调试模型时发现头数太少比如1或2模型表达能力明显不足生成描述很呆板头数太多比如16或32则训练不稳定容易过拟合对于图像描述这种任务4到8个头通常是甜点区。2.2 位置编码Positional Encoding给模型“空间感”注意力机制本身是“无序”的它只关心内容的相关性不知道“第一个词”和“第二个词”的顺序区别。这就像把一句话的单词全部打乱单靠自注意力很难恢复原句的语义。因此我们必须显式地告诉模型每个token的位置信息这就是位置编码的使命。Transformer原论文使用了一种非常巧妙的正弦余弦编码。它不是可学习的参数而是一个固定的、预先计算好的矩阵。对于序列中的第pos个位置在嵌入向量的第2i维使用正弦函数在第2i1维使用余弦函数频率随着维度i的增加而降低。这种设计使得模型能够轻松地学习到相对位置关系例如“隔了k个位置”这种模式。在实现PositionalEncoding类时关键就是正确构造这个[max_len, embed_dim]的矩阵然后在forward函数中将其加到输入的词嵌入向量上。记得最后要加一个Dropout这是提升模型泛化能力的标准操作。3. 动手实现从MultiHeadAttention到Transformer层理论聊得差不多了现在打开你的Jupyter Notebook我们开始写代码。作业的Transformer_Captioning.ipynb已经搭好了框架我们的任务就是填充那些TODO部分。3.1 实现MultiHeadAttention类这是第一个也是最重要的挑战。这个类要能通用地处理自注意力、掩码自注意力和交叉注意力。它的forward函数接收query, key, value和可选的attn_mask。实现步骤拆解投影与分头首先将输入的Q、K、V分别通过三个线性层self.key, self.query, self.value进行投影。然后把投影后的张量从(N, T, E)的形状重塑reshape为(N, T, h, D)再转置transpose为(N, h, T, D)其中h是头数D E/h是每个头的维度。这一步是把数据准备好让每个头能独立计算。计算注意力分数计算Q K.transpose(-1, -2)得到形状为(N, h, T, T)的原始分数矩阵。然后切记要除以sqrt(D)进行缩放。如果不做缩放当D较大时点积结果方差很大经过Softmax后权重会变得非常尖锐接近one-hot导致梯度变小训练困难。应用掩码如果提供如果传入了attn_mask比如在解码器掩码自注意力中需要屏蔽未来信息使用masked_fill方法将mask中为0的位置替换为一个极大的负值如-1e9这样在后续Softmax时这些位置的权重就会趋近于0。Softmax与Dropout对缩放并掩码后的分数矩阵在最后一个维度dim-1上应用Softmax得到注意力权重。接着立刻应用一个Dropoutself.attn_drop进行正则化这在训练时至关重要。加权聚合与合并将注意力权重与V相乘attn_weight V得到每个头的输出(N, h, T, D)。然后转置回(N, T, h, D)再重塑为(N, T, E)把多个头的输出在特征维度拼接起来。最终投影将拼接后的结果通过最后一个线性层self.proj映射回d_model维度得到最终的注意力输出。我调试这个模块时最容易出错的地方是张量形状的操作和转置的维度。一定要用print(x.shape)随时检查中间结果的形状是否符合预期。另一个坑是掩码的应用时机一定要在Softmax之前完成masked_fill。3.2 组装TransformerDecoderLayer有了多头注意力这个万能工具构建Transformer的解码器层就清晰多了。一个标准的Transformer解码器层包含三个子层掩码自注意力层Masked Self-Attention让当前生成的词只能关注到它之前的所有词包括自己这是保证生成过程自回归性的关键。交叉注意力层Cross-Attention这是连接视觉和语言的关键它的query来自解码器上一层的输出文本信息而key和value来自编码器的最终输出memory即图像特征。这让解码器在生成每一个词时都能“瞥一眼”图像确保描述内容不跑偏。前馈网络层Feed-Forward Network, FFN一个简单的两层MLP通常中间维度会扩大比如d_model512中间层dim_feedforward2048并带有ReLU激活和Dropout。它为每个位置的特征提供非线性变换。作业中采用了Pre-LNLayer Normalization的结构这与原论文的Post-LN不同。Pre-LN的意思是在每个子层注意力、FFN之前先做LayerNorm然后将子层的输出与子层的输入残差连接相加。公式大致是x x sublayer(LayerNorm(x))。这种结构现在更流行因为它能让训练更稳定梯度流动更好。在代码中你需要为每个子层实现这个“Norm - Sublayer - Dropout - Add”的流程。4. 构建完整的图像描述模型CaptioningTransformer现在我们把所有的零件组装成一辆能跑的“车”。对于图像描述任务我们通常使用一个编码器-解码器架构。但在这个作业的简化版本中图像特征已经由一个预训练好的CNN比如ResNet提取好了所以我们不需要实现视觉编码器只需要一个基于Transformer的解码器来生成文字。模型构建流程如下处理文本输入将单词索引captions通过一个嵌入层nn.Embedding转换为稠密向量然后加上位置编码使用我们刚实现的PositionalEncoding。处理图像输入将CNN提取的图像特征features形状(N, D)通过一个线性投影层self.visual_projection映射到与词向量相同的维度wordvec_dim并增加一个序列维度变成(N, 1, W)作为解码器的memory。你可以把这1理解为图像的一个“超级token”。生成因果掩码Causal Mask为了在训练时使用“教师强制”teacher forcing同时保证自回归性我们需要一个下三角矩阵作为掩码。torch.tril(torch.ones(T, T))可以生成一个主对角线及以下为1以上为0的矩阵。这样在预测第t个词时模型只能看到前t-1个词。通过Transformer解码器将加了位置编码的词向量作为tgt投影后的图像特征作为memory因果掩码作为tgt_mask一起送入我们搭建好的TransformerDecoder由多个TransformerDecoderLayer堆叠而成。输出词表分布将解码器的输出通过一个线性层self.output映射到词表大小vocab_size得到每个时间步上所有单词的得分logits。在sample推理函数中我们使用贪心解码从STARTtoken开始每一步都将当前已生成的部分序列输入模型取模型预测概率最高的词作为下一个词直到生成ENDtoken或达到最大长度。在实际产品中你可能会用束搜索Beam Search来获得更好的结果但贪心解码最简单也足以验证模型是否学会基本能力。5. 进阶实现视觉TransformerViT并思考其特性作业的后半部分引导我们实现一个简化版的ViTVision Transformer用于图像分类。这能让我们更深入地理解Transformer如何直接处理图像。5.1 实现PatchEmbeddingViT的核心思想是将图像视为一系列序列化的图像块Patch。PatchEmbedding层的工作就是将一张(N, C, H, W)的图片按照patch_size分割成(H/patch_size) * (W/patch_size)个不重叠的块。将每个块展平成一个向量维度是patch_size * patch_size * C。通过一个线性层self.proj将每个块的展平向量投影到固定的嵌入维度embed_dim。实现时巧妙运用reshape和permute来避免for循环是关键。例如对于形状为(N, C, H, W)的输入x目标输出是(N, num_patches, embed_dim)。你可以先将其变形为(N, C, num_patches_h, patch_size, num_patches_w, patch_size)然后调整维度顺序最后再展平。这一步如果写对了后面就一帆风顺。5.2 实现TransformerEncoderLayer并训练ViT的编码器层比解码器层更简单因为它只有自注意力层和前馈网络层没有交叉注意力。实现逻辑和之前的TransformerDecoderLayer中的自注意力子层、FFN子层几乎一样。将多个这样的编码器层堆叠起来前面加上PatchEmbedding和可学习的[CLS]token后面接一个分类头就构成了一个完整的ViT分类模型。在CIFAR-10这样的小数据集上训练ViT是个挑战。我最初的尝试用默认的小参数效果很差验证准确率不到40%。经过一番调参我发现了几点关键增大模型容量将embed_dim从128提升到768num_heads增加到12dim_feedforward增加到512。这给了模型更强的表达能力。调整学习率使用更小的学习率如1e-4配合AdamW优化器并加上weight_decay如1e-4防止过拟合。Batch Size的影响在总训练时长epoch数受限的情况下适当减小batch_size比如从64降到32有时能让模型在验证集上获得更好的性能这可能是因为小批量带来了更多的梯度更新次数和一定的正则化效果。经过这些调整模型在CIFAR-10上训练2个epoch后验证准确率可以从不到40%提升到接近48%。这虽然仍远低于CNN但验证了ViT架构的可行性。5.3 思考ViT在小数据集上为何不如CNN这是作业中一个非常好的思考题。根本原因在于归纳偏置Inductive Bias。CNN天生就内置了“局部性”和“平移不变性”的强假设这非常契合自然图像的结构——物体的特征通常由局部像素组合而成且无论物体在图像的哪个位置它还是那个物体。这种先验知识让CNN在数据量不足时也能学得很好。而ViT的注意力机制是全局的、内容驱动的它几乎没有关于图像空间的先验假设。它需要从数据中自己学习“相邻的像素块可能有关联”、“物体的形状是连续的”这些概念。这需要大量的数据。当数据量小时ViT就容易过拟合或者学不到有效的特征表示。如何改善ViT在小数据上的表现呢从我实战和阅读论文的经验看主要有几个方向更强的数据增强和正则化如RandAugment, MixUp, CutMix以及DropPath, Stochastic Depth等。知识蒸馏用一个在小数据上训练好的CNN教师模型来指导ViT学生模型训练把CNN的归纳偏置“迁移”过来。DeiT论文就用了这招。预训练-微调这是最有效的一招。先在超大规模数据集如ImageNet-21K, JFT上做监督或自监督预训练如MAE, DINO学到通用的视觉表征再到小数据集上做轻量级微调。混合架构与改进在ViT的早期引入卷积操作如卷积Patch Embedding或者使用层次化设计的Transformer如Swin Transformer显式地引入局部性先验。5.4 分析自注意力的计算成本最后我们来算算账。自注意力层的计算复杂度主要来自两部分计算注意力矩阵的O(N^2 * d)以及线性投影的O(N * d^2)。其中N是序列长度对于ViT就是图像块的数量d是隐藏维度。根据这个公式我们就能快速分析作业里的问题(i) 隐藏维度d加倍O(N^2 * d)项变为2倍O(N * d^2)项变为4倍。在ViT常见配置下N约几百d几百到上千后者往往占主导或相当所以总成本增加接近2到4倍。(ii) 图像高宽加倍N变为原来的4倍因为块数N (H*W) / P^2。O(N^2 * d)项暴增至16倍成为绝对主导总成本激增。(iii) 块大小P加倍N减少为原来的1/4。同理O(N^2 * d)项锐减至1/16总成本大幅下降。(iv) 层数加倍计算成本很直接就是原来的2倍。理解这些成本变化对于设计高效的模型至关重要。例如当你需要处理高分辨率图像时直接增大输入尺寸ii会带来巨大的计算负担此时增大块大小iii或采用分层注意力如Swin中的窗口注意力就是更明智的选择。写完所有这些代码并看着模型从最初的乱码开始逐渐能输出“a dog is running in the grass”这样像样的句子那种成就感是无与伦比的。这个过程里你不仅搞懂了Transformer的矩阵运算更理解了它如何让模型具备“全局视野”和“动态聚焦”的能力。这份从理论到实现的完整经历是只看论文或调用API无法替代的。希望你在实现过程中遇到bug时别灰心多打印中间变量形状多画图理解数据流动你一定能搞定它。