import pandas as pd
fig, axs = plt.subplots(3, 3, sharex=True, sharey=True, constrained_layout=True, figsize=(6, 6))
fig.subplots_adjust(hspace=0.1)
datasets = ['envi', 'enfr','enge']
types = ['enc', 'enc-dec' ,'dec']
Datasets = ['En-Vi', 'En-Fr','En-Ge']
def draw_map(ax,xlabels, ylabels, tt, attns, index): #画热力图
im = ax.imshow(attns,interpolation='none', cmap='Blues', vmin=0, vmax=1)
ax.set_title(tt, fontsize=10)
style='italic' )
return im
xlabels = '0 1 2 3 4 5'.split()
ylabels = '5 4 3 2 1 0'.split()
for i, data in enumerate(datasets):
datafile = '{}.csv'.format(data) #读取数据
data = pd.read_csv(datafile, delimiter='\t', header=None)
for j, tt in enumerate(types):
A = data.values[:,j*6:(j+1)*6] #取出所需热力图的数据
ax = axs[i, j]
im = draw_map(ax, xlabels,ylabels, '%s (%s)'%(Datasets[i],tt), A, j)
plt.xticks(range(len(xlabels)),xlabels) #设置横坐标label
plt.yticks(range(len(ylabels)),ylabels) #设置纵坐标label
cb_ax = fig.add_axes([1.0, 0.1, 0.02, 0.8]) #设置colarbar位置
cbar = fig.colorbar(im, cax=cb_ax) #共享colorbar
plt.tight_layout()
plt.show()
fig.savefig('weights_visualization.pdf', bbox_inches='tight', dpi=500)
画图结果：
