PyTorch模型的可解释性：注意力热图的可视化实现

【图书推荐】《PyTorch深度学习与计算机视觉实践》-CSDN博客

基于Vision Transformer的mini_ImageNet图片分类实战_imagenet数据集-CSDN博客

对于注意力模型来说，它与传统的卷积神经网络或者循环神经网络的一个最显著的区别就是使用了注意力机制，使得模型可以聚集在某个特定的位置上。而如何可视化地获取到注意力集中的点就是一项令人兴奋的任务。本小节将结合已完成的Vision Transformer模型，实现Vision Transformer注意力热图的可视化。在开始之前首先需要导入对应的库包：

pip install grad-cam

下面开始实现注意力热图的可视化。

1. 载入模型

首先要做的是对模型的重用。在这里可以直接使用训练好的模型，并完成模型参数的载入，代码如下：

import vison_transformer,get_dataset
device = "cuda" if torch.cuda.is_available() else "cpu"
model = vison_transformer.VisionTransformer(depth = 6)
save_path = "./saver/vit_model.pth"
model.load_state_dict(torch.load(save_path))
model = model.cuda().eval()		#注意这里强制使用模型的推断模式

需要说明，对于任意一个模型，包括现在想可视化的这个Vision Transformer模型，其可视化的目标都是“一个层”的参数，而这个层需要由层名称获取。我们可以使用print(model)打印模型中所有层名称，如图11-8所示。

图11-8  打印所有层名称

此时可以看到，blocks是模型的一个一级层名称，norm1则是blocks中一个block的最低级的层，这里需要的就是此层名称，完整的名称为“model.blocks[-1].norm1”。它在GradCAM类的初始化中使用。

2. 输入图像

这里使用训练集中的第一幅图进行图像的输入。由于PyTorch模型在推断时也要保证图像的输入格式，因此使用如下代码输入图像：

image = (get_dataset.train_image_list[10357])/255.
input_tensor = torch.tensor(image).float()	#将数据格式转换为tensor
input_tensor = torch.permute(input_tensor, [2, 0, 1])	#调整维度
input_tensor = input_tensor.unsqueeze(0)	#扩充一个数据维度
input_tensor = input_tensor.to(device)		#将数据发送到硬件GPU中

3. 使用可视化库包函数

接下来使用可视化库包完成图像热点的可视化，代码如下：

from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image, preprocess_image
#是对层获取的参数维度进行变更，在下面GradCAM类的初始化中使用
#height与width分别是获取到的图片大小
def reshape_transform(tensor, height=84, width=84):
    #去掉cls token，并调整维度，将图形大小调整成初始大小
    result = tensor[:, 1:, :].reshape(tensor.size(0),3,height, width)
    #将通道维度放到第一个位置
    result = torch.permute(result, [0, 2, 3,1])
    return result

#GradCAM是可视化的初始类，target_layers针对的是VIT模型中可视化输出的那个层，使用层名称进行输出
cam = GradCAM(
model=model,			#初始化了目标模型
target_layers=[model.blocks[-1].norm1],#提取的目标层，一定要和模型中的层名称对应
use_cuda=True,							#是否使用GPU加速
reshape_transform=reshape_transform)	#对提取的层进行整形的方法

#计算 grad-cam
target_category = None #可以指定一个类别，或者使用 None 表示最高概率的类别
grayscale_cam = cam(input_tensor=input_tensor, targets=target_category)
grayscale_cam = grayscale_cam[0, :]
#将 grad-cam 的输出叠加到原始图像上
visualization = show_cam_on_image(image, grayscale_cam)
#保存可视化结果
cv2.cvtColor(visualization, cv2.COLOR_RGB2BGR, visualization)
cv2.imwrite('cam.jpg', visualization)

这里的GradCAM是初始化的类，代码中对目标模型、提取的目标层以及对应的整形方法都做了显式说明。需要注意，target_layer需要结合前期目标中的那个层名称，才可以获得对应的参数。可视化展示效果如图11-9所示。

图11-9  对Vision Transformer模型进行热力图提取的可视化展示

可以很明显地看到，当获取到model.blocks[-1].norm1这个层的参数后，可视化表示的注意力模型会重点关注鸟类的头部、腰部以及脚部，而model.blocks[-2].norm2则侧重鸟类的头部。同时，通过对其他层的展示也可以看到，随着层级的不同，观察的重点也不相同。有兴趣的读者可以自行比较和查看。

《PyTorch深度学习与计算机视觉实践（人工智能技术丛书）》(王晓华)【摘要 书评 试读】- 京东图书 (jd.com)