1. 项目概述：为什么从可视化入手理解Transformer？

如果你在深度学习领域待过一段时间，尤其是涉足自然语言处理（NLP）或者计算机视觉（CV），那么“Transformer”这个词对你来说，可能既熟悉又陌生。熟悉的是，它几乎成了所有SOTA（State-of-the-Art）模型的基石，从BERT、GPT系列到ViT、Swin Transformer，无处不在。陌生的是，当你翻开那篇著名的论文《Attention Is All You Need》，面对那一堆矩阵运算、多头注意力机制和位置编码公式时，是不是感觉头大如斗，仿佛在看天书？我最初接触Transformer时，也有同感。公式推导和代码实现固然重要，但对于建立直观理解，它们有时反而成了障碍。

这就是为什么我坚信，从可视化开始，是理解Transformer最高效的路径。我们的大脑天生对图像和动态过程更敏感。与其在抽象的数学符号里打转，不如先看看这个模型内部到底在“看”什么、“想”什么。一个清晰的可视化，能瞬间将复杂的“黑箱”操作，转化为我们可以直观感知的信息流和注意力模式。这就像给你一张复杂电路板的X光透视图，你一眼就能看清电流的走向和元件的连接，远比死记硬背电路图要来得深刻。

本文的目标读者很明确：所有被Transformer的数学细节劝退，但又渴望理解其核心思想的开发者、学生和爱好者。无论你是想复现一个模型，还是仅仅为了在面试或技术讨论中能言之有物，通过可视化建立直观认知，都是你绕不开的第一步。接下来，我不会堆砌公式，而是带你一起，通过几个关键的可视化场景，亲手“拆解”Transformer，看看自注意力机制如何工作，编码器-解码器如何交互，以及位置信息如何被巧妙编码。我们会使用一些现成的工具和少量代码，让你不仅能看懂，还能自己动手复现这些可视化过程。

2. 核心思路：拆解Transformer的可视化维度

在动手之前，我们需要一个清晰的“地图”。Transformer模型虽然结构统一，但内部信息流动复杂，从不同角度可视化，能看到不同的风景。盲目地看所有细节，只会让人更加混乱。因此，我把Transformer的可视化分解为四个核心维度，由表及里，逐步深入。

2.1 维度一：模型整体架构流图

这是最宏观的视角，目的是看清Transformer的“骨架”。我们熟悉的编码器-解码器结构、残差连接、层归一化（Add & Norm）模块是如何组织在一起的？数据（比如一个句子）从输入到输出，究竟流经了哪些部件？

为什么先看这个？因为它建立了全局认知。就像在组装乐高之前，你得先看说明书上的成品图。这个维度的可视化通常是一个静态的、高层次的框图。虽然论文里的图1已经给出了经典结构，但我们可以用更生动的工具（如TensorBoard的add_graph功能，或Netron这样的模型结构查看器）来加载一个实际的Transformer模型（比如Hugging Face的bert-base-uncased），直观地看到每一个子模块的层级关系和数据流向。这能帮你确认：“哦，原来一个编码器层里，先经过多头注意力，再加残差和归一化，然后再过前馈网络，再来一次残差和归一化。”

2.2 维度二：注意力权重的热力图

这是Transformer的灵魂所在，也是可视化最能大显身手的地方。自注意力机制的核心是计算一个序列中所有元素（例如句子中的所有单词）两两之间的关联程度，这个关联程度就是注意力权重（Attention Weights），最终形成一个权重矩阵。

可视化形式：热力图（Heatmap）。假设我们输入句子“The cat sat on the mat”。在某个注意力头（Head）中，模型在计算“sat”这个词的表示时，会去看句子中所有其他词（包括“sat”自己）的权重。我们可以将这个权重向量（长度为句子长度）画成一个热力图的一行。将所有词的注意力行堆叠起来，就得到一个方形的热力图。图中颜色越亮（如红色、黄色），表示权重越高，关联越强。

看什么？

1. 对角线模式：模型是否高度关注词语本身？这在某些情况下是合理的。
2. 语法/语义关联：“cat”和“sat”之间是否有强连接？（主谓关系）“on”和“mat”之间呢？（介词-宾语关系）
3. 多头差异：不同的注意力头可能专注于不同的关系模式。有的头可能专看局部语法（相邻词），有的头可能专看长程依赖（句首和句尾的词）。通过对比不同头的热力图，你能直观感受到“多头”设计的精妙。

2.3 维度三：词嵌入向量的空间分布

在进入Transformer之前，词语会被转换成高维向量（词嵌入）。经过Transformer各层处理后，这些向量的含义会不断被上下文信息所丰富和修正。我们可以通过降维技术（如t-SNE或UMAP），将这些高维向量投影到2D或3D空间，观察它们的分布变化。

可视化形式：2D/3D散点图。我们可以对比输入层的词向量、经过第一层编码器后的向量、以及经过最后一层编码器后的向量在空间中的分布。

看什么？

1. 语义聚类：语义相近的词（如“cat”, “dog”）是否在空间中靠得更近？经过网络深层处理后，这种聚类是否更加清晰？
2. 句法角色分离：名词、动词、介词是否会形成不同的簇？
3. 上下文影响：同一个词在不同句子中（如“bank”在“river bank”和“bank account”中），其最终的向量表示是否会漂移到不同的区域？这直观展示了Transformer如何根据上下文动态调整词义。

2.4 维度四：训练过程中的动态变化

模型不是静态的。观察训练过程中上述可视化的动态变化，能让你理解模型是如何“学习”的。例如，注意力热力图从最初的随机、混乱模式，逐渐收敛到有意义的语法/语义模式；词向量在空间中的分布从一团混沌，逐渐形成有结构的聚类。

可视化工具：这通常需要结合训练日志和可视化库（如Matplotlib动画，或TensorBoard的嵌入投影仪Embedding Projector）来实现。虽然实现稍复杂，但带来的认知提升是巨大的。你能亲眼见证一个“婴儿”模型如何通过数据喂养，逐渐成长为能理解语言的“成人”。

注意：在实际操作中，我们不必一次性实现所有维度的可视化。对于初学者，我强烈建议从注意力权重热力图开始，因为它最直接地揭示了Transformer的核心计算过程，且实现相对简单，效果冲击力强。掌握了它，你就抓住了理解Transformer的“牛鼻子”。

3. 实操准备：环境、工具与一个小型案例

理论说得再多，不如亲手运行一遍。为了让可视化过程顺畅，我们需要搭建一个轻量级的环境，并选择一个合适的“小白鼠”模型和数据集。这里的原则是：最小化环境依赖，最大化直观效果。

3.1 环境与工具链选择

我们不搞复杂的分布式训练环境，一个能跑通PyTorch或TensorFlow的Python环境足矣。我个人的偏好是PyTorch，因为其动态图特性在调试和可视化时更灵活。

核心工具包清单：

• 深度学习框架：PyTorch(>=1.9) 或TensorFlow(>=2.4)。本文示例将以PyTorch为主。
• Transformer模型库：Hugging Face Transformers。这是我们的“军火库”，提供了数以千计预训练的Transformer模型及其易用的接口。安装命令：pip install transformers。
• 可视化库：
  • Matplotlib&Seaborn：用于绘制静态的热力图、曲线图等。Seaborn能让热力图更美观。
  • Plotly或Bokeh：可选，用于制作交互式图表，体验更佳。
  • scikit-learn：主要用于其中的TSNE或PCA降维算法，用于词向量可视化。
• 数据处理：pandas,numpy用于基础数据操作。

一个避坑点：不同版本的库可能存在API差异。建议使用虚拟环境（如conda或venv）管理依赖，并尽量安装较新但稳定的版本。例如，可以创建一个conda环境：conda create -n transformer-viz python=3.9，然后在该环境中安装上述包。

3.2 选择一个合适的模型与数据

对于初次可视化，目标不是解决复杂任务，而是看清机制。因此，我们选择一个小型、结构清晰、且预训练好的模型。

模型选择：DistilBERT为什么不直接用最大的BERT或GPT？因为它们层数多、参数量大，注意力头也多（BERT-base有12层*12头=144个注意力矩阵），初看容易眼花缭乱。DistilBERT是BERT的一个蒸馏版本，它保留了BERT 97%的性能，但只有6层编码器，参数量减少40%。层数少意味着我们需要可视化的矩阵更少，更容易跟踪信息流，同时其注意力模式依然具有代表性。从Hugging Face加载它非常简单：

from transformers import DistilBertModel, DistilBertTokenizer model = DistilBertModel.from_pretrained('distilbert-base-uncased') tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased') model.eval() # 设置为评估模式，关闭dropout等训练层

数据选择：一个简短的例句我们需要一个足够短、语法语义清晰的句子，这样热力图不会太大，便于观察。同时，句子最好能包含一些有趣的语法关系。我选择这个经典例句：“The quick brown fox jumps over the lazy dog.”这个句子包含了所有字母，且具有清晰的主谓宾结构（fox jumps over dog）和修饰关系（quick brown, lazy）。长度适中，非常适合演示。

3.3 实现注意力权重提取的“钩子”技巧

预训练模型通常不会直接返回中间层的注意力权重。我们需要在PyTorch中使用前向钩子（Forward Hook）来“拦截”模型在前向传播过程中产生的这些中间变量。这是实现可视化的关键技术点。

原理：PyTorch的nn.Module允许我们注册一个钩子函数，这个函数会在该模块的前向传播计算完成后被调用，并可以获取该模块的输入和输出。

操作步骤：

1. 定位注意力层：在DistilBert中，注意力计算发生在DistilBertLayer里的attention模块中。
2. 定义钩子函数：这个函数接收模块本身、其输入和输出三个参数。我们从输出中提取注意力权重（通常是一个元组或字典，需要查看模型源码确认结构）。
3. 注册钩子：将钩子函数注册到我们感兴趣的特定层（比如第0层和第5层，看首尾差异）。
4. 运行模型：将处理好的句子输入模型，触发前向传播，钩子函数会自动执行，并将截获的权重存储到我们定义的全局列表中。

一个简易的钩子示例框架：

attentions = [] # 用于存储各层注意力权重的列表 def hook_fn(module, input, output): # 假设output是一个元组 (hidden_states, attention_weights) # 具体结构需要根据实际模型调整，查阅文档或源码是关键 if isinstance(output, tuple) and len(output) > 1: attentions.append(output[1].detach().cpu()) # 取注意力权重，并移出计算图 # 注册钩子到特定层 layer_index = 0 # 例如，第一层 model.transformer.layer[layer_index].attention.register_forward_hook(hook_fn)

实操心得：不同模型（BERT, GPT, ViT）输出注意力权重的格式可能不同。最可靠的方法是直接查阅该模型类的源代码（通常在Hugging Face Transformers库的modeling_xxx.py文件中），找到注意力计算返回的具体内容。这是调试过程中最关键的一步，也是深入理解模型实现的好机会。

4. 核心可视化实现：绘制注意力热力图

环境备好，钩子设下，数据就位，现在让我们来生成第一张，也是最重要的一张图——注意力热力图。我们将以DistilBERT处理示例句子为例，展示从数据准备到成图的完整流程。

4.1 数据预处理与模型前向传播

首先，我们需要将文本句子转化为模型能理解的输入ID张量，并准备好注意力掩码（Attention Mask，用于区分真实词和填充词）。

sentence = “The quick brown fox jumps over the lazy dog.” inputs = tokenizer(sentence, return_tensors=“pt”) # 返回PyTorch张量 # inputs 是一个字典，包含 ‘input_ids‘ 和 ’attention_mask‘ with torch.no_grad(): # 禁用梯度计算，节省内存和计算资源 outputs = model(**inputs, output_attentions=True) # 关键：要求模型输出注意力权重 # outputs.attentions 现在是一个元组，包含了每一层每一个头的注意力权重 # 形状通常是 (层数, 批大小, 头数, 序列长度, 序列长度)

这里有一个关键参数：output_attentions=True。对于许多Hugging Face模型，直接通过这个参数就能获取所有层的注意力权重，比手动挂钩子更方便。但了解钩子技术仍然重要，因为它更灵活，可以获取任何中间变量。

4.2 提取并解析注意力权重张量

假设我们通过上述方式获得了all_attentions。它是一个包含6个元素的元组（对应DistilBERT的6层）。每个元素是一个张量，形状为[1, 12, 11, 11]（批大小1，12个注意力头，序列长度11，序列长度11）。

import numpy as np layer_attentions = all_attentions[0] # 取第一层的注意力权重 # layer_attentions.shape: torch.Size([1, 12, 11, 11]) # 我们关心第一个样本，第一个注意力头 head_idx = 0 attention_weights = layer_attentions[0, head_idx].numpy() # 转换为numpy数组 # attention_weights 是一个 11x11 的矩阵

这个11x11的矩阵M，其中M[i, j]就表示在生成第i个词的上下文向量时，模型分配给第j个词的注意力分数。分数经过softmax归一化，所以每一行的和都为1。

4.3 使用Seaborn绘制热力图

现在，我们将这个矩阵用热力图呈现出来。为了可读性，我们需要将矩阵的索引（0,1,2...）替换为对应的单词。

import seaborn as sns import matplotlib.pyplot as plt # 获取分词后的tokens（包含[CLS]和[SEP]） tokens = tokenizer.convert_ids_to_tokens(inputs[‘input_ids’][0]) # tokens: [‘[CLS]’, ‘the’, ‘quick’, ‘brown’, ‘fox’, ‘jumps’, ‘over’, ‘the’, ‘lazy’, ‘dog’, ‘.’, ‘[SEP]’] # 创建图形 fig, ax = plt.subplots(figsize=(10, 8)) # 绘制热力图 sns.heatmap(attention_weights, xticklabels=tokens, yticklabels=tokens, cmap=“Reds”, # 红色系，越红权重越高 square=True, # 保持单元格为正方形 cbar_kws={“shrink”: 0.8}, # 调整颜色条大小 ax=ax) ax.set_title(f“DistilBERT Layer 0 - Head {head_idx+1} Attention Weights”) ax.set_xlabel(“Key Tokens (被关注的词)”) ax.set_ylabel(“Query Tokens (正在计算的词)”) plt.xticks(rotation=45, ha=“right”) # 旋转x轴标签，防止重叠 plt.tight_layout() plt.show()

4.4 结果分析与模式解读

运行上述代码，你会得到一张11x11的热力图。让我们来仔细“阅读”这张图：

1. 观察对角线：你可能会发现对角线上的颜色比较亮。这很常见，表示模型在计算某个词的表示时，会给予该词本身一定的关注（自注意力）。但这并非绝对，特别是在深层网络，模型可能更关注其他有语义关联的词。

2. 寻找语法关系：

   • 找到“jumps”所在的行（Query）。看这一行中，哪些Key的格子最亮？你很可能会发现“fox”的格子非常亮。这完美捕捉了“fox jumps”的主谓关系。
   • 再看“over”所在的行，它可能同时关注“jumps”和“dog”，捕捉了“jumps over”的动词短语和“over the dog”的介词短语结构。
   • “dog”所在的行，可能会强烈关注“the”和“lazy”，捕捉了定冠词和形容词修饰。

3. 比较不同注意力头：将head_idx从0改到11，重复绘图。你会发现不同的头关注点截然不同。有的头可能专门关注“下一个词”（类似RNN），有的头关注“句法中心词”，有的头可能关注标点符号。这就是“多头”的魅力——它允许模型同时从不同的表示子空间学习不同的关系。

4. 比较不同层：将layer_attentions = all_attentions[0]改为all_attentions[5]（最后一层）。你会发现深层的注意力模式往往比浅层更加稀疏和专业化。浅层的注意力可能更偏向于局部语法和词性，而深层的注意力可能更聚焦于高层次的语义关联和任务相关特征。

注意事项：注意力权重并不直接等同于“重要性”或“语义关联度”。它是模型为了优化最终任务（如掩码语言建模）而学习到的一种计算策略。有时，一些看似不相关的词之间出现较高权重，可能是模型学习到的一种间接的、服务于全局目标的关联模式。解读时需要结合具体任务上下文。

5. 进阶可视化：词向量空间与训练动态

掌握了注意力热力图，你已经理解了Transformer最核心的机制。接下来，我们可以探索另外两个维度的可视化，它们能帮助你从表示学习和模型演进的角度获得更完整的认知。

5.1 词向量降维与可视化

词向量是模型对语言的“内部理解”。我们可以提取不同网络层输出的词向量（即hidden_states），通过降维技术将其投射到二维平面进行观察。

步骤：

1. 提取各层隐藏状态：在调用模型时，使用output_hidden_states=True参数。outputs.hidden_states会是一个元组，包含嵌入层和每一层编码器输出的隐藏状态。
2. 选择目标层和词：例如，我们关注名词“fox”、“dog”和动词“jumps”、“over”在嵌入层、第3层、第6层（最后一层）的向量表示。
3. 应用降维：使用sklearn.manifold.TSNE进行降维。t-SNE擅长在低维空间保持高维数据的局部结构，适合可视化聚类。
4. 绘制散点图：为不同层的向量使用不同颜色或标记，观察同一个词在不同层的向量位置如何变化。

from sklearn.manifold import TSNE import matplotlib.pyplot as plt # 假设 hidden_states 是包含各层输出的元组 emb_layer_vec = hidden_states[0][0, [4, 9, 5, 6]] # [CLS], fox, dog, jumps, over 的索引 layer3_vec = hidden_states[3][0, [4, 9, 5, 6]] layer6_vec = hidden_states[6][0, [4, 9, 5, 6]] # 合并所有向量 all_vecs = torch.cat([emb_layer_vec, layer3_vec, layer6_vec], dim=0).numpy() # 创建标签 labels = [‘fox(emb)’, ‘dog(emb)’, ‘jumps(emb)’, ‘over(emb)’, ‘fox(L3)’, ‘dog(L3)’, ‘jumps(L3)’, ‘over(L3)’, ‘fox(L6)’, ‘dog(L6)’, ‘jumps(L6)’, ‘over(L6)’] # t-SNE降维 tsne = TSNE(n_components=2, perplexity=3, random_state=42) # perplexity调小以适应少量样本 vecs_2d = tsne.fit_transform(all_vecs) # 绘图 plt.figure(figsize=(10, 8)) colors = [‘red’, ‘blue’, ‘green’, ‘purple’] for i, (label, (x, y)) in enumerate(zip(labels, vecs_2d)): layer = label.split(‘(’)[1][:-1] # 提取emb, L3, L6 color = colors[i % 4] # 同一个词在不同层用同色 marker = ‘o’ if ‘emb’ in layer else ‘s’ if ‘L3’ in layer else ‘^’ # 用不同标记区分层 plt.scatter(x, y, c=color, marker=marker, s=100, label=label if i<4 else “”) plt.annotate(label, (x, y), xytext=(5, 2), textcoords=‘offset points’) plt.legend([‘Embedding Layer’, ‘Layer 3’, ‘Layer 6’]) plt.title(‘Word Vector Trajectory through Transformer Layers (t-SNE)’) plt.xlabel(‘t-SNE Dimension 1’) plt.ylabel(‘t-SNE Dimension 2’) plt.grid(True, alpha=0.3) plt.show()

解读：你可能会观察到，“fox”和“dog”在嵌入层可能距离较远，但随着层数加深，因为它们都是句子中的动物名词（宾语），它们的向量在空间中可能会逐渐靠近。而“jumps”和“over”作为动词和介词，也可能形成自己的簇。这直观展示了Transformer层如何根据上下文信息，动态调整和聚合词义表示。

5.2 跟踪训练过程中的注意力演化

要观察训练动态，你需要在一个小任务上（如文本分类）从头或继续训练一个Transformer模型，并在每个训练周期（Epoch）保存特定层、特定头的注意力矩阵。然后，你可以将这些矩阵制作成动画。

简化实现思路：

1. 定义训练循环，并在每个epoch结束后，用验证集的一个固定样例进行推理，并保存其注意力权重。
2. 存储数据：将每个epoch的注意力矩阵（例如layer0_head0）保存为一个列表或字典。
3. 制作动画：使用matplotlib.animation.FuncAnimation，将每个epoch的热力图依次显示。

import matplotlib.animation as animation from matplotlib.animation import PillowWriter # 假设 attention_by_epoch 是一个列表，包含每个epoch的注意力矩阵 fig, ax = plt.subplots(figsize=(8, 6)) def update(frame): ax.clear() attn = attention_by_epoch[frame] sns.heatmap(attn, ax=ax, cbar=False, cmap=“Reds”, square=True) ax.set_title(f“Epoch {frame+1}”) ax.set_xticklabels(tokens) ax.set_yticklabels(tokens) plt.xticks(rotation=45) ani = animation.FuncAnimation(fig, update, frames=len(attention_by_epoch), interval=500) # 保存为GIF ani.save(“attention_evolution.gif”, writer=PillowWriter(fps=2))

解读动画：在训练初期，注意力矩阵可能近乎均匀或呈现随机模式。随着训练进行，你会看到清晰的模式逐渐“浮现”和“稳定”。例如，主谓、动宾之间的连线会变得越来越亮。这个过程生动地展示了模型如何从数据中学习到有意义的语言结构。

6. 常见问题与排查技巧实录

在实际操作可视化过程中，你几乎一定会遇到各种问题。下面是我在多次实践中总结的一些典型问题及其解决方案，希望能帮你节省大量调试时间。

6.1 注意力权重全为零或NaN

现象：绘制出的热力图一片漆黑，或者控制台提示权重包含NaN值。可能原因与排查：

1. 模型处于训练模式：如果模型有Dropout层，在训练模式下，注意力权重可能会被随机置零。务必在提取权重前调用model.eval()。
2. 梯度计算干扰：在前向传播时没有使用torch.no_grad()上下文管理器，导致计算图保存，可能在某些情况下引发异常。确保推理代码被with torch.no_grad():包裹。
3. 数据包含padding：如果句子经过批量处理且长度不一，注意力掩码（attention_mask）会屏蔽掉padding位置。提取权重时，这些被屏蔽位置的权重可能是0或一个极小的负数（softmax前）。这是正常的。可视化时，你可以选择只关注非padding部分。
4. 模型输出结构理解错误：这是最常见的问题。outputs.attentions的结构因模型而异。可能是元组，也可能是张量。使用print(type(outputs.attentions))和print(outputs.attentions[0].shape)来仔细检查其形状。参考对应模型的文档或源代码（如modeling_distilbert.py中的forward函数返回值）是唯一正解。

6.2 热力图可读性差

现象：图像模糊，颜色对比不明显，或者标签重叠看不清。优化技巧：

1. 调整色彩映射（cmap）：‘viridis’,‘plasma’适合连续数据，‘Reds’,‘Blues’适合突出高值。使用sns.color_palette(“rocket”, as_cmap=True)可以获得更专业的配色。
2. 聚焦特定区域：如果序列很长（如512），整个热力图会非常密集。可以只可视化前50个token的注意力，或者只可视化某个特定词（如“[CLS]” token）对所有其他词的注意力（即矩阵的一行）。
3. 格式化标签：对于长序列，可以间隔显示标签，或者将标签旋转90度。使用plt.xticks(rotation=90, ha=‘center’)。
4. 增加图形尺寸：通过fig, ax = plt.subplots(figsize=(16, 12))增大画布。
5. 对权重取对数：有时注意力分布非常尖锐（大部分接近0，少数几个接近1）。可以对权重取np.log(weight + 1e-10)后再绘图，以增强低权重的可视对比度。

6.3 词向量可视化点堆叠在一起

现象：t-SNE图上的点全部挤在一团，看不出任何结构。排查与调整：

1. 检查向量是否已归一化：如果向量模长差异巨大，t-SNE可能无法很好工作。尝试对向量进行L2归一化：vec = vec / np.linalg.norm(vec, axis=1, keepdims=True)。
2. 调整t-SNE参数：
   • perplexity（困惑度）：这是最重要的参数。它大致表示每个点考虑多少近邻。对于少量点（如几十个），应设置一个较小的值（5-20）；对于大量点，使用默认值30或更高。需要反复尝试。
   • learning_rate（学习率）：通常介于10到1000。如果点形成一团团紧凑的球，尝试降低学习率；如果点散乱无章，尝试提高学习率。
   • n_iter（迭代次数）：确保迭代足够（通常>1000），让算法有足够时间收敛。
3. 尝试其他降维方法：t-SNE的结果具有随机性，且不保证全局结构。可以同时尝试PCA（主成分分析）作为对比。PCA是线性方法，能保持全局方差，有时能提供互补的视角。
4. 样本量是否过少？t-SNE需要一定数量的点才能学习到流形结构。如果只可视化几个词，可能难以形成有意义的布局。可以尝试可视化一个迷你批次（batch）中所有句子的所有词。

6.4 钩子函数未触发或获取数据不对

现象：注册了钩子，但列表attentions仍然是空的，或者获取的数据形状不对。解决步骤：

1. 确认注册时机：钩子必须在模型前向传播之前注册。确保你的register_forward_hook调用在model(inputs)之前。
2. 确认目标模块：使用print(model)或print(model.transformer.layer[0])来查看模型的确切结构，确保你注册钩子的路径是正确的。模块名可能因模型版本不同而变化。
3. 在钩子函数内打印：在钩子函数的第一行添加print(“Hook called!”, type(output))，并检查是否被打印。如果没有，说明钩子未正确触发。
4. 深入探查输出结构：在钩子函数内，使用print(output)或if isinstance(output, tuple): print(len(output))来查看输出到底是什么。它可能是一个张量，也可能是一个复杂的嵌套结构。你需要根据实际情况调整提取权重的代码。
5. 考虑使用nn.Module的_modules属性：对于更复杂的嵌套结构，可能需要递归地查找和注册钩子。

可视化Transformer不是一个一蹴而就的任务，而是一个迭代探索的过程。从一张简单的热力图开始，逐步增加复杂度，对比不同层、不同头、不同输入，你的直觉会在这个过程中被慢慢培养起来。当你能看着这些彩色方块，脑海中自动映射出模型理解语言的逻辑时，你就真正从“知道”Transformer，变成了“看懂”Transformer。