importnumpyasnpimportmatplotlib.pyplotaspltfrompylabimportmplmpl.rcParams['font.sans-serif']= ['STZhongsong']mpl.rcParams['axes.unicode_minus']=Truedefheatmap(data,row_labels,col_labels,ax=None,cbar_kw={},cbarlabel="",**kwargs):ifnotax:ax=plt.gca()im=ax.imshow(data,**kwargs)cbar=ax.figure.colorbar(im,ax=ax,**cbar_kw)cbar.ax.set_ylabel(cbarlabel,rotation=-90,va="bottom")ax.set_xticks(np.arange(data.shape[1]))ax.set_yticks(np.arange(data.shape[0]))ax.set_xticklabels(col_labels)ax.set_yticklabels(row_labels)ax.tick_params(top=True,bottom=False,labeltop=True,labelbottom=False)plt.setp(ax.get_xticklabels(),rotation=-20,ha="right",rotation_mode="anchor")foredge,spineinax.spines.items():spine.set_visible(False)ax.set_xticks(np.arange(data.shape[1]+1)-.5,minor=True)ax.set_yticks(np.arange(data.shape[0]+1)-.5,minor=True)ax.grid(which="minor",color="w",linestyle='-',linewidth=3)ax.tick_params(which="minor",bottom=False,left=False)returnim,cbarif__name__=="__main__":x,y=list("你这是干什么啦？"),list("你要吃什么？")s=np.random.random([len(x),len(y)])fig,ax=plt.subplots()im,cbar=heatmap(s,x,y,ax=ax,cmap="YlGn",cbarlabel="attention scores")fig.tight_layout()plt.savefig("test",dpi=300,bbox_inches='tight')plt.show()