import torch from d2l import torch as d2l def show_heatmaps(matrices, xlabel, ylabel, titles=None, figsize=(2.5, 2.5),cmap='Reds'): d2l.use_svg_display() num_rows, num_cols = matrices.shape[0], matrices.shape[1] fig, axes = d2l.plt.subplots(num_rows, num_cols, figsize=figsize, sharex=True, sharey=True, squeeze=False) for i, (row_axes, row_matrices) in enumerate(zip(axes, matrices)): for j, (ax, matrix) in enumerate(zip(row_axes, row_matrices)): pcm = ax.imshow(matrix.detach().numpy(), cmap=cmap) if i == num_rows - 1: ax.set_xlabel(xlabel) if j == 0: ax.set_ylabel(ylabel) if titles: ax.set_title(titles[j]) fig.colorbar(pcm, ax=axes, shrink=0.6); attention_weights = torch.eye(10).reshape((1, 1, 10, 10)) show_heatmaps(attention_weights, xlabel=’Keys’, ylabel=’Queries’)

这段代码调用了之前定义的`show_heatmaps`函数，并传入了一个注意力权重矩阵作为参数。注意力权重矩阵是一个10x10的单位矩阵，表示了一个查询和键之间的关注程度。

在这段代码中，我们首先创建了一个10x10的单位矩阵作为注意力权重。然后，我们调用了`show_heatmaps`函数，并将注意力权重矩阵作为参数传递给它。同时，我们还指定了x轴和y轴上的标签为"Keys"和"Queries"。

`show_heatmaps`函数会将注意力权重矩阵可视化为一个热图，并在图像中显示x轴和y轴上的标签。由于我们没有提供任何标题，所以图像中不会显示标题。

如果你运行这段代码，你将看到一个显示了注意力权重矩阵的热图，其中x轴和y轴上分别标有"Keys"和"Queries"。

希望对你有所帮助！如果还有其他问题，请随时提问。