1.简介

最近读了不少与注意力机制相关的文章，对网络进行了一定的修改，想通过神经网络注意力热图来看看效果，看了不少博文，总结一下这几天实验的结果。

通过阅读，大部分博文都是通过得到神经网络的输出特征图（大部分是最后一个卷积层的输出，不过其实任意层特征图都可以，看你想观察那个部分），将特征图resize到输入图像的大小，通过cv2中的函数叠加到原图像就可以了。感觉很简单，但其实中间还是有很多细节存在的。我主要用的方法是CAM类激活图的方法。

?2.类激活图（CAM）

CAM方法是文章Learning Deep Features for Discriminative Localization提出来的。

该文章的主要内容就是上面的图片了，

1)首先得到一个输出特征图[B, C, H, W]，B为Batch size，如果是输入一张图片则B=1；C为最后一个卷积层的输出通道数；H和W为最后一个卷积层输出特征图的宽高。

2)然后得到你训练好的模型的分类头的权重classifier.weight。注意，这里分类层的输入通道数要与输出特征图的输出通道数匹配。

3)将每个通道的特征图进行加权求和，权重就为第二步的权重。最后的输出结果就为每一类的注意力热图了。

3.CAMpytorch代码（步骤已做注释）

1）首先加载需要用到的库，定义类别数，实例化网络，加载预训练模型，再生成一个只到最后一个卷积层的网络model_features，得到分类层的权重fc_weights，生成分类任务对应的字典，都进入eval（）模式。

import os
import numpy as np
import cv2
import torch
import torch.nn as nn
from torchvision import transforms, datasets
from PIL import Image

#加载自己的网络
from model import model

class_num = 5

model_ft = model(num_classes=class_num)
model_ft.load_state_dict(torch.load('pretrain.pth', map_location=lambda  storage, loc:storage))

model_features = nn.Sequential(*list(model_ft.children())[:-2])
fc_weights = model_ft.state_dict()['classifier.weight'].cpu().numpy()  #numpy数组取维度fc_weights[0].shape->(5,960)
class_ = {0:'car', 1:'bird', 2:'tree', 3:'sky', 4:'person'}
model_ft.eval()
model_features.eval()

2）设置一个输入图片的预处理方式，导入图片，图片预处理，将图片输入到model_features中得到输出特征图[B, C, 7, 7]。

若是还要在图片上写入预测的结果，就将图片再输入到model_ft中，得到预测出的tensor，shape为[1,5]，再通过softmax得到每一类的预测概率，将概率和对应的索引从大到小排列，打印出每一类的概率。

data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
        "val": transforms.Compose([transforms.Resize(256),
                                   transforms.CenterCrop(224),
                                   transforms.ToTensor(),
                                   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}

img_path = '/data/test.jpg'             #单张测试
_, img_name = os.path.split(img_path)
features_blobs = []
img = Image.open(img_path).convert('RGB')
img_tensor = data_transform['val'](img).unsqueeze(0) #[1,3,224,224]
features = model_features(img_tensor).detach().cpu().numpy()  #[1,960,7,7]
# print(features.shape)
logit = model_ft(img_tensor)  #[1,2] -> [ 3.3207, -2.9495]
h_x = torch.nn.functional.softmax(logit, dim=1).data.squeeze()  #tensor([0.9981, 0.0019])

probs, idx = h_x.sort(0, True)      #按概率从大到小排列
probs = probs.cpu().numpy()  #if tensor([0.0019,0.9981]) ->[0.9981, 0.0019]
idx = idx.cpu().numpy()  #[1, 0]
# print(probs.shape)
# print(class_)
for i in range(class_num):
    print('{:.3f} -> {}'.format(probs[i], class_[idx[i]]))  #打印预测结果

3）写一个函数，输入为最后一层特征图、分类头的权重、概率从大到小的类别索引，输出每一类对应的注意力热图集。

def returnCAM(feature_conv, weight_softmax, class_idx):
    bz, nc, h, w = feature_conv.shape        #1,960,7,7
    output_cam = []
    for idx in class_idx:  #只输出预测概率最大值结果不需要for循环
        feature_conv = feature_conv.reshape((nc, h*w))  # [960,7*7]
        cam = weight_softmax[idx].dot(feature_conv.reshape((nc, h*w)))  #(5, 960) * (960, 7*7) -> (5, 7*7) （n,）是一个数组，既不是行向量也不是列向量
        cam = cam.reshape(h, w)
        cam_img = (cam - cam.min()) / (cam.max() - cam.min())  #Normalize
        cam_img = np.uint8(255 * cam_img)                      #Format as CV_8UC1 (as applyColorMap required)

        #output_cam.append(cv2.resize(cam_img, size_upsample))  # Resize as image size
        output_cam.append(cam_img)
    return output_cam
# CAMs = returnCAM(features, fc_weights, [idx[0]])  #输出预测概率最大的特征图集对应的CAM
CAMs = returnCAM(features, fc_weights, idx)  #输出预测概率最大的特征图集对应的CAM
print(img_name + ' output for the top1 prediction: %s' % class_[idx[0]])

4）得到featmap后，读入输入图片的尺寸，将featmap的尺寸resize为原图大小，通过cv2.applyColorMap函数得到更直观的注意力热图，再与原图叠加起来，再将之前输出的类别概率写到图上，最后将featmap保存下来。

img = cv2.imread(img_path)
height, width, _ = img.shape  #读取输入图片的尺寸
heatmap = cv2.applyColorMap(cv2.resize(CAMs[0], (width, height)), cv2.COLORMAP_JET)  #CAM resize match input image size
result = heatmap * 0.3 + img * 0.5    #比例可以自己调节


text = '%s %.2f%%' % (class_[idx[0]], probs[0]*100) 				 #激活图结果上的文字显示
cv2.putText(result, text, (210, 40), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.9,
            color=(123, 222, 238), thickness=2, lineType=cv2.LINE_AA)
CAM_RESULT_PATH = r'/data/heatmap/'   #CAM结果的存储地址
if not os.path.exists(CAM_RESULT_PATH):
    os.mkdir(CAM_RESULT_PATH)
image_name_ = img_name.split(".")[-2]
cv2.imwrite(CAM_RESULT_PATH + image_name_ + '_' + 'pred_vit' + class_[idx[0]] + '.jpg', result)  #写入存储磁盘