### 如何用Python实现注意力机制热力图 (Attention Map) 为了展示注意力机制中的权重分布情况，绘制注意力热力图是一种直观的方式。下面将以 PyTorch 实现 Transformer 的 Multi-Head Attention 为基础来创建注意力热力图。 #### 准备工作 确保安装了必要的库： ```bash pip install torch matplotlib seaborn pandas ``` #### 创建简单的多头自注意力模块并获取注意力分数 定义一个多头注意层，在前向传播过程中保存注意力得分矩阵以便后续可视化[^3]。 ```python import torch.nn as nn import math class ScaledDotProductAttention(nn.Module): """Scaled dot-product attention mechanism.""" def __init__(self, temperature, attn_dropout=0.1): super().__init__() self.temperature = temperature self.dropout = nn.Dropout(attn_dropout) self.softmax = nn.Softmax(dim=2) def forward(self, q, k, v, mask=None): # Compute the scaled dot product. attn = torch.bmm(q, k.transpose(1, 2)) attn = attn / self.temperature if mask is not None: attn = attn.masked_fill(mask == 0, -1e9) attn = self.softmax(attn) attn = self.dropout(attn) output = torch.bmm(attn, v) return output, attn # 返回output的同时返回attn用于绘图 ``` #### 绘制注意力热力图函数 利用 `matplotlib` 和 `seaborn` 库来进行图形化表示: ```python import numpy as np import seaborn as sns; sns.set_theme() import matplotlib.pyplot as plt def plot_attention_map(att_mat, tokens=None, save_path='attention.png'): num_heads = att_mat.shape[0] fig, axes = plt.subplots(num_heads//2, 2, figsize=(8, 6*num_heads/4)) for i in range(num_heads): ax = axes[i // 2][i % 2] sns.heatmap(att_mat[i], annot=True, fmt=".1f", cmap="YlGnBu", square=True, cbar=False, ax=ax) if tokens: ax.set_xticklabels(tokens, rotation=45, ha='right') ax.set_yticklabels(tokens[::-1]) ax.set_title(f'Head {i+1}') ax.xaxis.tick_top() plt.tight_layout() plt.savefig(save_path) plt.show() ``` 此代码片段允许传入由上述自定义的 `ScaledDotProductAttention` 类产生的注意力张量以及可选参数——token 列表作为输入，从而生成每个多头注意力子层对应的热力图，并将其保存到指定路径下。 通过这种方式，可以清晰地观察不同头部之间的关注模式差异及其在整个序列上的分布特性。