有时候需要可视化观察一下模型中间输出的结果,写个博客简单记录一下,方便后面直接粘贴复用
import numpy as np
import matplotlib.pyplot as plt
#xq_out.shape [B,H,W]
normalized_tensor = (xq_out[0] - xq_out[0].min()) / (xq_out[0].max() - xq_out[0].min())
plt.imshow(normalized_tensor.cpu(), cmap='hot')
plt.colorbar()
plt.savefig('./heatmap_org2.png', bbox_inches='tight', pad_inches=0) #, dpi=600 调整分辨率