Torch 论文复现：梯度加权类激活映射 Grad-CAM

荷碧TongZJ

已于 2023-05-02 22:35:10 修改

阅读量4k

点赞数 9

分类专栏：计算机视觉文章标签： pytorch 人工智能 python 卷积神经网络

于 2022-07-14 00:01:11 首次发布

本文链接：https://blog.csdn.net/qq_55745968/article/details/125773860

版权

计算机视觉专栏收录该内容

15 篇文章 37 订阅

订阅专栏

Grad-CAM 全称 Gradient-weighted Class Activation Mapping，用于卷积神经网络的可视化，甚至可以用于语义分割

不过我是主要研究目标检测的，在看论文的时候就没有在意语义分割的部分

Grad-CAM 的前身是 CAM，CAM 的基本的思想是求分类网络某一类别得分对高维特征图 (卷积层的输出) 的偏导数，从而可以该高维特征图每个通道对该类别得分的权值；而高维特征图的激活信息 (正值) 又代表了卷积神经网络的所感兴趣的信息，加权后使用热力图呈现得到 CAM

在分类网络的可视化方面上，Grad-CAM 并没有太大的进步，两者都是基于全局平均池化的结构进行手动的推导，针对权值的计算进行改进

先讲一下 Grad-CAM 论文的基本思路，再讲讲我的 idea

以经典的分类网络 ResNet50 为例，卷积层 (下采样 4 次) + 全局平均池化 + 全连接层

当输入图像的 shape 为 [3, 224, 224] 时：

卷积层：[3, 224, 224] -> [2048, 14, 14]，记输出为 A
全局平均池化：[2048, 14, 14] -> [2048, 1, 1]，记输出为 F
全连接层：[2048, 1, 1] -> [2048, ] -> [1000, ]，对应 1000 个类别的得分，记输出为 Y

对于第 class 个类别，其得分为 $gif.latex?Y%5E%7Bclass%7D$ (重在思路，以下计算忽略线性层的偏置)：

$gif.latex?Y%5E%7Bclass%7D%3D%5Csum_%7Bch%3D1%7D%5E%7B2048%7Dw_%7Bch%7D%5E%7Bclass%7D%5Ccdot%20F_%7Bch%7D$

$gif.latex?F_%7Bch%7D%20%3D%20%5Cfrac%7B1%7D%7B14%5E2%7D%20%5Csum_%7Bx%3D1%7D%5E%7B14%7D%20%5Csum_%7By%3D1%7D%5E%7B14%7D%20A_%7Bx%5C%20y%7D%5E%7Bch%7D$

类别得分 $gif.latex?Y%5E%7Bclass%7D$ 对高维特征图每一个元素 $gif.latex?A_%7Bx%5C%20y%7D%5E%7Bch%7D$ 的偏导数为：

$gif.latex?%5Cfrac%7B%5Cpartial%20Y%5E%7Bclass%7D%7D%7B%5Cpartial%20A_%7Bx%5C%20y%7D%5E%7Bch%7D%7D%20%3D%20%5Cfrac%7B1%7D%7B14%5E2%7Dw_%7Bch%7D%5E%7Bclass%7D$

因为使用了全局平均池化，所以最终求得偏导数与类别、高维特征图的通道有关，而与高维特征图的像素位置无关

Grad-CAM 表示为：

$gif.latex?M_%7Bx%5C%20y%7D%5E%7Bclass%7D%20%3D%20ReLU%28%5Csum_%7Bch%3D1%7D%5E%7B2048%7D%5Cfrac%7B%5Cpartial%20Y%5E%7Bclass%7D%7D%7B%5Cpartial%20A_%7Bx%5C%20y%7D%5E%7Bch%7D%7D%5Ccdot%20A_%7Bx%5C%20y%7D%5E%7Bch%7D%29$

从等式中可以看出，Grad-CAM 其实就是求解了高维特征图对某类别得分的权值 (即贡献率)，并在通道维度上对高维特征图进行加权，以表征每一个位置对该类别得分的贡献程度 (上采样之后拓展到全图)

对于更复杂的网络，其偏导数的推理肯定更为复杂，而“与高维特征图的像素位置无关”这个结论也将因为全局平均池化层的移除而不适用

在 Torch 框架中，我们可以借助反向传播梯度的机制，对复杂网络的 Grad-CAM 进行求解 (源代码在文末)

因为这个 Grad-CAM 具有非常好的定位性能，所以论文作者又对 Guided Backprop 下了手

Guided Backprop 是以图像 x 为自变量 (即 requires_grad)，以某一类别的得分为 Loss 进行梯度的反向传播，最终以图像的梯度作为 Guided Backprop (梯度表征了该像素点对类别得分的贡献)

但是 Guided Backprop 的定位效果很差，从最上面的那幅图可以看到，当设置类别为“猫”时，Guided Backprop 中“狗”的轮廓特别的清晰

Grad-CAM 是一个掩膜矩阵，Guided Backprop 是图像的梯度矩阵，两者利用广播机制进行相乘，即得到 Guided Grad-CAM

Grad-CAM 复现

参考代码：https://github.com/jacobgil/pytorch-grad-cam

上面这份代码可使用 pip install grad-cam 下载；Grad-CAM++ 的论文我没有看，不知道是不是这篇论文的代码

我看了这份代码，总结出绘制分类网络的 Grad-CAM 的流程如下：

在目标层 (通常为最后一层卷积层) 设置 hook，在前向传播时保存该层输出张量 (高维特征图)，在反向传播时保存该层输出张量的梯度 (高维特征图的梯度，用于与高维特征图加权，生成 Grad-CAM)
以类别得分作为 Loss 反向传播梯度，对高维特征图的梯度进行处理 (在该代码中则在对梯度进行全局平均池化，其实我觉得使用逐元素的梯度更好、更通用)，然后使用梯度对高维特征图进行加权；而后将加权的高维特征图沿通道求和得到 Grad-CAM，使用 ReLU 函数剔除负值，再使用双线性插值 (上采样) 使 Grad-CAM 的尺寸与原图像相同
对于多个目标层，会产生多个 Grad-CAM，此时使用求平均的方法进行聚合 (但是对于 YOLO 这种路径聚合型的网络，更需要的是多个 Grad-CAM 进行叠加，也就是使用求最值的方法进行聚合)

而绘制 Guided Grad-CAM 的流程大有不同：

重新定义 ReLU 的反向传播机制 (我测试过了，如果没有这一步最后的效果会很差)
以类别得分作为 Loss 反向传播梯度，记录图像的梯度矩阵作为 Guided Backprop
Grad-CAM 是一个掩膜矩阵，Guided Backprop 是图像的梯度矩阵，两者利用广播机制进行相乘，即得到 Guided Grad-CAM

我对这份代码有几个比较不满意的点：

对于高维特征图没有使用逐元素的梯度，通用性不强 (没法用到我计划研究的 YOLO)
多个 Grad-CAM 的聚合方式单一 (只有求平均)，只可用于防抖动，不可用于叠加
封装程度太高，对于调包小伙比较方便，对于二创玩家来说很头疼
有些代码的写法明显太复杂了，有失优雅

我仿照这份代码的基本流程，进行了重建：

from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn


def normalize(x, eps=1e-7, **kwargs):
    ''' 归一化'''
    x -= x.min(**kwargs)
    x /= x.max(**kwargs) + eps
    return x


class BP_ReLU(nn.Module):

    def forward(self, x):
        return self.Func.apply(x)

    class Func(torch.autograd.Function):

        @staticmethod
        def forward(ctx, x):
            y = F.relu(x)
            ctx.save_for_backward(x, y)
            return y

        @staticmethod
        def backward(ctx, gy):
            x, y = ctx.saved_tensors
            gx = F.relu(gy) * (x > 0)
            return gx


class Grad_CAM:
    ''' Gradient-weighted Class Activation Mapping
        model: 卷积神经网络模型
        target_layers: 可视化的目标层列表
        agg_fun: 多个目标层的 CAM 聚合方式 (max, avg)
        act_replace: 更换激活函数'''

    def __init__(self, model, target_layers, agg_fun='avg',
                 act_replace=[(nn.ReLU, BP_ReLU)]):
        self.model = model.eval()
        # 注册挂钩, 以保存目标层输出张量及其梯度
        self.handle = sum([[
            ty.register_forward_hook(
                lambda module, x, y: self.activation.append(y.detach())
            ),
            ty.register_full_backward_hook(
                lambda module, gx, gy: self.grad.append(gy[0])
            )] for ty in target_layers], [])
        # 生成聚合函数
        kwargs = dict(dim=1, keepdims=True)
        self.agg_fun = {'max': lambda fm: fm.max(**kwargs)[0],
                        'avg': lambda fm: fm.mean(**kwargs)}[agg_fun]
        # 记录需要被更换的网络层
        self.act_replace = act_replace

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        for handle in self.handle: handle.remove()

    def __call__(self, file_list, transform, file_mode='Project/%s.jpg'):
        ''' file_list: 图像的路径列表
            transform: array(BGR) -> tensor 的转换函数'''
        folder = Path(file_mode).parent
        folder.mkdir(exist_ok=True)
        # 保存 BGR 图像, 绘制热力图时使用
        bgr_list = []
        for file in file_list:
            img = cv.imread(file)
            assert isinstance(img, np.ndarray), f'OpenCV can\'t open file: {file}'
            bgr_list.append(img)
        # 转换成 tensor, 并在 batch 维度上拼接
        x = torch.stack([transform(bgr) for bgr in bgr_list], dim=0)
        grad_cam = self.get_cam(x)
        guided_bp = self.get_bp(x)
        # 可视化并存储结果
        for file, bgr, bp, cam in zip(file_list, bgr_list, guided_bp, grad_cam):
            stem = Path(file).stem
            # 根据 grad-cam 生成热力图
            cam_mask = np.uint8(np.round(cam * 255)).repeat(3, -1)
            heat_map = cv.applyColorMap(cam_mask, cv.COLORMAP_JET)
            # Grad-CAM
            cam_image = cv.addWeighted(bgr, 0.5, heat_map, 0.5, 0)
            cv.imwrite(file_mode % f'{stem}.Grad-CAM', cam_image)
            # Guided-Backprop
            gbp_image = self.write_grad(bp)
            cv.imwrite(file_mode % f'{stem}.Guided-Backprop', gbp_image)
            # Guided Grad-CAM
            ggc_image = self.write_grad(bp * cam_mask)
            cv.imwrite(file_mode % f'{stem}.Guided Grad-CAM', ggc_image)

    def target(self, y, *args, **kwargs):
        print(y.shape)
        raise NotImplementedError('No maximization goal is defined')

    def replace(self, old, new, model=None):
        ''' 更换模型中的 ReLU 模块'''
        model = self.model if model is None else model
        for key, module in model._modules.items():
            if isinstance(module, old):
                model._modules[key] = new()
            self.replace(old, new, module)

    def get_cam(self, x):
        ''' Grad-CAM'''
        # 对激活函数进行还原, 撤销更换
        for new, old in self.act_replace:
            self.replace(old, new)
        self.model.zero_grad()
        size = x.shape[-2:]
        # 清空挂钩读取的激活图和梯度
        self.activation, self.grad = [], []
        # 前向传播, 反向传播
        tar = self.target(self.model(x))
        tar.backward()
        # 特征图根据反向传播的梯度进行组合, 沿通道求和得到该目标层的 CAM
        # 将多个目标层的 CAM 上采样后, 量化后再沿通道维度拼接, 使用聚合函数合并
        grad_cam = self.agg_fun(torch.cat([F.interpolate(
            normalize(F.relu(act * grad).sum(dim=1, keepdims=True)),
            size=size, mode='bilinear', align_corners=False
        ) for act, grad in zip(self.activation, reversed(self.grad))], dim=1))
        return grad_cam.permute(0, 2, 3, 1).data.numpy()

    def get_bp(self, x):
        ''' Guided Backprop'''
        # 对激活函数进行更换, 以优化最终的可视化结果
        for old, new in self.act_replace:
            self.replace(old, new)
        self.model.zero_grad()
        x.requires_grad = True
        # 前向传播, 反向传播
        tar = self.target(self.model(x))
        tar.backward()
        guided_bp = x.grad
        return guided_bp.permute(0, 2, 3, 1).data.numpy()[..., ::-1]

    def write_grad(self, img, eps=1e-7):
        img -= np.mean(img)
        img /= (np.std(img) + eps)
        img = img * 0.1
        img = img + 0.5
        img = np.clip(img, 0, 1)
        return np.uint8(img * 255)

使用 Grad_CAM 对象的步骤如下：

重写 Grad_CAM 的 target 方法，在其中定义 Loss 的计算 (因为这个 Loss 是我们想要最大化的，与平常的最小化不同，所以我给这个函数起名 target)
加载卷积神经网络模型，与目标层一同实例化 Grad_CAM 管理器 (关注 __init__ 方法)，再把想要可视化的图像文件名、BGR 图像变换为 tensor 的函数名传进管理器 (关注 __call__ 方法)，在当前目录的 Project 文件夹中找到可视化结果

因为 YOLOv6、v7 还没有开始研究，所以只做了分类网络的实验 (ResNet50)

from torchvision import models
import torchvision.transforms as tf
from PIL import Image
import cv2 as cv


def image_tran(file_or_bgr, shape=None):
    transform = [tf.Resize(shape) if shape else tf.Compose([]),
                 tf.ToTensor(),
                 tf.Normalize(mean=[0.485, 0.456, 0.406],
                              std=[0.229, 0.224, 0.225])]
    # 如果是文件, 则读取
    if isinstance(file_or_bgr, str):
        file_or_bgr = Image.open(file_or_bgr)
    else:
        file_or_bgr = file_or_bgr[..., ::-1].copy()
    return tf.Compose(transform)(file_or_bgr)


class My_CAM(Grad_CAM):

    def target(self, y, *args, **kwargs):
        # 对于分类网络, 取每一张图片类别分数的最大值之和
        max_score = y.max(dim=1)
        return max_score[0].sum()


parent = Path(r'D:\Information\Python\Laboratory\data')
img_file = ['both.png', 'dog_cat.jfif', 'dogs.png']

model = models.resnet50(pretrained=True)
target_layers = [model.layer4]

with My_CAM(model, target_layers) as cam:
    for img in img_file:
        cam([str(parent / img)], image_tran)

以下四幅图分别为：原图、Guided Backprop、Grad-CAM、Guided Grad-CAM

在梯度图中，亮度高于背景的像素处为正梯度 (白色区)，亮度低于背景的像素处为负梯度 (黑色区)

正梯度与负梯度的交界区 (或是正梯度的集中区) 表征了对“狗”这个类别有正增益的轮廓，而负梯度的集中区则表征了其它类别的轮廓