注意图

import matplotlib.pyplot as plt

def visualize_attention(attention_map, original_image):
    plt.imshow(original_image[0], cmap='gray')
    plt.title('原始图像')
    plt.show()

    plt.imshow(attention_map[0, 0].detach().cpu().numpy(), cmap='hot', interpolation="nearest")
    plt.title('注意力图')
    plt.show()

# 假设你可以访问名为 'attention_layer' 的 ModulatedAttLayer 实例
# 以及一个名为 'input_image' 的输入张量

output, feature_maps = attention_layer(input_image)
visualize_attention(feature_maps[1], input_image)