保存整个图形很简单，只需使用savefig函数：

f.savefig('filename.png')

有许多文件格式可以保存，这些通常可以从文件名的扩展名中正确推断出来。有关详细信息，请参阅文档。
savefig函数接受一个参数bbox_inches，它定义了要保存的图形区域。要将单个子图保存到文件中，您可以使用子图Axes对象的边界框来计算相应的值。
把它们放在一起，你的代码看起来像这样：

f, ax = plt.subplots(1, 2)
for i, img in enumerate([img1, img2]):    
    grads = # my visualization codes
    # visualize grads as heatmap
    ax[i].imshow(grads, cmap='jet')

    # Save the subplot.
    bbox = ax[i].get_tightbbox(f.canvas.get_renderer())
    f.savefig("subplot{}.png".format(i),
              bbox_inches=bbox.transformed(f.dpi_scale_trans.inverted()))

# Save the whole figure.
f.savefig("whole_figure.png")