简明代码介绍类激活图CAM, GradCAM, GradCAM++

  类激活图（class activation map, CAM)能够显示输入图像各区域对于分类神经网络指定类别提供信息的多少，可以帮助我们更好的理解神经网络的工作过程。关于CAM网上讲的有不少，我这里重点给出CAM的简明代码实现，并讨论一下各种情况下CAM的效果。CAM已经发展出了好几种方法，有CAM, GradCAM, GradCAM++, SmoothGradCAM++, ScoreCAM, SSCAM, ISCAM等。这里只介绍前三种，其他暂时不打算学。网上给出的代码要么完整而太长，要么太短而不完整，都不易快速理解，我这里给出完整又简短的写法，比较好理解。如果对其中涉及到的hook钩子方法不太熟悉，可以看我另一篇博客专门介绍钩子方法。

一、CAM

  CAM仅适用于网络结构是最后一个卷积层+GAP(即global average pooling,全局平均池化)+fc的情况。这种情况时，最后一个卷积层输出的特征图feature_map（在下面代码中我们命名为hook_a）经过GAP池化处理后会从C×H×W变成C×1×1，这样最后的全连接层fc的输入节点数就是C，所以fc的权重weights就是C × num_classes的。如果我们把weights中每个类别的向量和最后卷积层的特征图相乘后再相加，就可以得到一个H×W尺寸的图，这个图成为热图heatmap，它反映了feature_map中每个像素对最后分类的重要程度。而卷积神经网络始终只是进行缩放而没有旋转变换，所以hook_a和原输入图片的坐标位置是对应的，因此可以resize后和原图片叠加起来用于显示原图中每个区域对指定类的贡献程度。
  好吧，用人类语言解释这个事太费劲了，估计大家没怎么看懂，show you the code：

import numpy as np
import torch.nn.functional as F
import torchvision.models as models
from torchvision.transforms.functional import normalize, resize, to_tensor, to_pil_image
import matplotlib.pyplot as plt
from matplotlib import cm
from PIL import Image

net = models.resnet18(pretrained=True).cuda()

hook_a = None
def _hook_a(module,inp,out):
    global hook_a
    hook_a = out

submodule_dict = dict(net.named_modules())
target_layer = submodule_dict['layer4']
hook1 = target_layer.register_forward_hook(_hook_a)

img_path = 'images/border_collie2.jpg'
img = Image.open(img_path, mode='r').convert('RGB')
img_tensor = normalize(to_tensor(resize(img, (224, 224))),
                           [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]).cuda()

scores = net(img_tensor.unsqueeze(0))
hook1.remove()
class_idx = 232 # class 232 corresponding to the border collie
weights = net.fc.weight.data[class_idx,:]
cam = (weights.view(*weights.shape, 1, 1) * hook_a.squeeze(0)).sum(0)
cam = F.relu(cam)
cam.sub_(cam.flatten(start_dim=-2).min(-1).values.unsqueeze(-1).unsqueeze(-1))
cam.div_(cam.flatten(start_dim=-2).max(-1).values.unsqueeze(-1).unsqueeze(-1))
cam = cam.data.cpu().numpy()

heatmap = to_pil_image(cam, mode='F')
overlay = heatmap.resize(img.size, resample=Image.BICUBIC)
cmap = cm.get_cmap('jet')
overlay = (255 * cmap(np.asarray(overlay) ** 2)[:, :, :3]).astype(np.uint8)
alpha = .7
result = (alpha * np.asarray(img) + (1 - alpha) * overlay).astype(np.uint8)
plt.imshow(result)

图1.原图

图2. CAM法输出结果

  我们可以看出来，ResNet18并没有关注到图中所有的边牧，它只关注了最像的一个，这就足以让它做出分类判断了。

二、GradCAM

  与CAM相比，GradCAM不再使用全连接层的权重，而是巧妙计算了指定类输出得分相对于卷积层特征图的梯度，不再要求最后一个卷积层和全连接层之间必须有GAP层，适用范围更广。

图3. GradCAM法核心公式

  
  用论文中这个公式来解释，其中 $A_{ij}^k$表示卷积层的输出特征图，k是通道数，i,j是像素数，$y^c$表示第c类的输出向量，$\frac{\partial y^c}{\partial A_{ij}^k}$表示的就是第c类对特征图的梯度。对此梯度按每通道求平均，也可以得到一个k维向量${\alpha }_k^c$, 就是代码中的weights，用它来取代CAM中的weights也可以进行计算。对于适用了末层卷积层和全连接层之间使用了GAP的情况下，用GradCAM法和CAM法计算的结果是相同的。下面还是给出完整代码：

import numpy as np
import torch.nn.functional as F
import torchvision.models as models
from torchvision.transforms.functional import normalize, resize, to_tensor, to_pil_image
import matplotlib.pyplot as plt
from matplotlib import cm
from PIL import Image

net = models.resnet18(pretrained=True).cuda()

hook_a = None
def _hook_a(module,inp,out):
    global hook_a
    hook_a = out

hook_g = None
def _hook_g(module,inp,out):
    global hook_g
    hook_g = out[0]

submodule_dict = dict(net.named_modules())
target_layer = submodule_dict['layer4']

hook1 = target_layer.register_forward_hook(_hook_a)
hook2 = target_layer.register_backward_hook(_hook_g)

img_path = 'images/border_collie3.jpg'
img = Image.open(img_path, mode='r').convert('RGB')
img_tensor = normalize(to_tensor(resize(img, (224, 224))),
                           [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]).cuda()

scores = net(img_tensor.unsqueeze(0))
class_idx = 232 # class 232 corresponding to the border collie
loss = scores[:,class_idx].sum()
loss.backward()
hook1.remove()
hook2.remove()

weights = hook_g.squeeze(0).mean(dim=(1,2))
cam = (weights.view(*weights.shape, 1, 1) * hook_a.squeeze(0)).sum(0)
cam = F.relu(cam)
cam.sub_(cam.flatten(start_dim=-2).min(-1).values.unsqueeze(-1).unsqueeze(-1))
cam.div_(cam.flatten(start_dim=-2).max(-1).values.unsqueeze(-1).unsqueeze(-1))
cam = cam.data.cpu().numpy()

heatmap = to_pil_image(cam, mode='F')
overlay = heatmap.resize(img.size, resample=Image.BICUBIC)
cmap = cm.get_cmap('jet')
overlay = (255 * cmap(np.asarray(overlay) ** 2)[:, :, :3]).astype(np.uint8)
alpha = .7
result = (alpha * np.asarray(img) + (1 - alpha) * overlay).astype(np.uint8)
plt.imshow(result)

  注意再GradCAM中使用的loss就是最后输出的指定类的得分，使用这个loss反向传播，而这里提取的梯度也不是常规情况下我们关注的权重的梯度，而是特征图的梯度。
  理论上GradCAM法可以用任意卷积层的特征图和特征图的梯度计算这个cam图：

图4. GradCAM法使用ResNet18不同层画出的cam图

  但是从各层画出的结果中可以看出，只有最后层画出的cam图有意义，这个原因我还没有细想。

三、GradCAM++

  GradCAM++为了效果更好，引入了二阶梯度和三阶梯度，在代码中计算的时候为了方便，分别用一阶梯度的平方和三次方来替代。代码和GradCAM基本相同，只需要把计算weights的一句替换为下面代码即可：

##weights = hook_g.squeeze(0).mean(dim=(1,2))
grad_2 = hook_g.pow(2)
grad_3 = grad_2 * hook_g
denom = 2 * grad_2 + (grad_3 * hook_a).sum(dim=(2, 3), keepdim=True)
nan_mask = grad_2 > 0
grad_2[nan_mask].div_(denom[nan_mask])
weights = grad_2.squeeze_(0).mul_(torch.relu(hook_g.squeeze(0))).sum(dim=(1, 2))

图5. GradCAM++法输出结果