show_heatmaps-未完待续

1. 代码

# -*- coding: utf-8 -*-
# @Project: zc
# @Author: ZhangChu
# @File name: show_heatmaps
# @Create time: 2021/11/6 7:30

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



def show_heatmaps(matries, xlabel, ylabel, tiles=None, figsize=(2.5, 2.5), cmap='Reds'):
	"""

	:param matries: matries = [要显示的行数，要显示的列数，查询的数目，键的数目]
	:param xlabel:  X 轴表示的标签含义
	:param ylabel:  Y 轴表示的标签含义
	:param tiles:
	:param figsize:
	:param cmap:
	:return:
	"""
	d2l.use_svg_display()
	num_rows, num_cols = matries.shape[0], matries.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, matries)):
		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 tiles:
				ax.set_titles(tiles[j])

	fig.colorbar(pcm, ax=axes, shrink=0.6)


attention_weight = torch.eye(10).reshape((1, 1, 10, 10))
show_heatmaps(attention_weight, xlabel='key', ylabel='Queries')

plt.show()