注意力机制热力图代码

以下是一个简单的注意力机制热力图的代码示例，使用了Python和PyTorch库：

import torch
import torch.nn as nn
import matplotlib.pyplot as plt

# 定义注意力机制热力图函数
def plot_attention(attention):
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(1, 1, 1)
    ax.imshow(attention, cmap='hot')
    ax.set_xticks([])
    ax.set_yticks([])
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['bottom'].set_visible(False)
    ax.spines['left'].set_visible(False)
    plt.show()

# 生成示例注意力权重
attention_weights = torch.tensor([[0.2, 0.3, 0.5],
                                  [0.1, 0.6, 0.3],
                                  [0.4, 0.2, 0.4]])

# 绘制热力图
plot_attention(attention_weights)

在这个示例中，`attention_weights`是一个3x3的注意力权重矩阵。通过调用`plot_attention`函数，我们可以将注意力机制的热力图显示出来。热力图使用热图颜色映射来表示不同位置的注意力权重值。

请注意，这只是一个简单的示例，实际应用中的注意力机制可能更加复杂。你可以根据自己的需求来调整和修改代码。希望对你有帮助！