Pytorch复现导向反向传播Guided Backpropagation

前言

  笔者在学习Grad-Cam算法对应的论文时，注意到该论文利用导向反向传播Guided Backpropagation来可视化细粒度信息，用Grad-Cam来定位判别性区域，大致如下图所示。因此，便对导向反向传播Guided Backpropagation算法进行了学习与复现。如果你想了解Grad-Cam算法的复现，可参考这篇博客：https://editor.csdn.net/md/?articleId=129650976

一、导向反向传播Guided Backpropagation的原理

关于导向反向传播Guided Backpropagation的原理可参考下图。

二、导向反向传播Guided Backpropagation的复现

注意，只需修改下段代码中输入图像的路径及输出图像的路径，然后运行即可。

import os
import cv2
import torch
from torch import nn
from PIL import Image
from torchvision import models
from torchvision import transforms

class Guided_backprop():
    def __init__(self, model):
        self.model = model
        self.image_reconstruction = None
        self.activation_maps = []
        self.model.eval()
        self.register_hooks()

    def register_hooks(self):
        def first_layer_hook_fn(module, grad_in, grad_out):
            # 在全局变量中保存输入图片的梯度，该梯度由第一层卷积层
            # 反向传播得到，因此该函数需绑定第一个 Conv2d Layer
            self.image_reconstruction = grad_in[0]

        def forward_hook_fn(module, input, output):
            # 在全局变量中保存 ReLU 层的前向传播输出
            # 用于将来做 guided backpropagation
            self.activation_maps.append(output)

        def backward_hook_fn(module, grad_in, grad_out):
            # ReLU 层反向传播时，用其正向传播的输出作为 guide
            # 反向传播和正向传播相反，先从后面传起
            grad = self.activation_maps.pop()
            # ReLU 正向传播的输出要么大于0，要么等于0，
            # 大于 0 的部分，梯度为1，
            # 等于0的部分，梯度还是 0
            grad[grad > 0] = 1

            # grad_in[0] 表示 feature 的梯度，只保留大于 0 的部分
            positive_grad_in = torch.clamp(grad_in[0], min=0.0)
            # 创建新的输入端梯度
            new_grad_in = positive_grad_in * grad

            # ReLU 不含 parameter，输入端梯度是一个只有一个元素的 tuple
            return (new_grad_in,)

        # 获取 module，这里只针对 alexnet，如果是别的，则需修改
        modules = list(self.model.features.named_children())

        # 遍历所有 module，对 ReLU 注册 forward hook 和 backward hook
        for name, module in modules:
            if isinstance(module, nn.ReLU):
                module.register_forward_hook(forward_hook_fn)
                module.register_backward_hook(backward_hook_fn)

        # 对第1层卷积层注册 hook
        first_layer = modules[0][1]
        first_layer.register_backward_hook(first_layer_hook_fn)

    def visualize(self, input_image, target_class):
        # 获取输出，之前注册的 forward hook 开始起作用
        model_output = self.model(input_image)
        self.model.zero_grad()
        pred_class = model_output.argmax().item()

        # 生成目标类 one-hot 向量，作为反向传播的起点
        grad_target_map = torch.zeros(model_output.shape,
                                      dtype=torch.float)
        if target_class is not None:
            grad_target_map[0][target_class] = 1
        else:
            grad_target_map[0][pred_class] = 1

        # 反向传播，之前注册的 backward hook 开始起作用
        model_output.backward(grad_target_map)
        # 得到 target class 对输入图片的梯度，转换成图片格式
        result = self.image_reconstruction.data[0].permute(1, 2, 0)
        return result.numpy()


def normalize(I):
    # 归一化梯度map，先归一化到 mean=0 std=1
    norm = (I - I.mean()) / I.std()
    # 把 std 重置为 0.1，让梯度map中的数值尽可能接近 0
    norm = norm * 0.1
    # 均值加 0.5，保证大部分的梯度值为正
    norm = norm + 0.5
    # 把 0，1 以外的梯度值分别设置为 0 和 1
    norm = norm.clip(0, 1)
    return norm


if __name__ == '__main__':
    I = Image.open('./test_0002_aligned.jpg').convert('RGB')
    transform = transforms.Compose([
        transforms.Resize((224,224)),
        transforms.CenterCrop((224,224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    tensor = transform(I).unsqueeze(0).requires_grad_()
    model = models.alexnet(pretrained=True)
    guided_bp = Guided_backprop(model)
    result = guided_bp.visualize(tensor, None)
    result = normalize(result)
    result= result[:, :, ::-1]*255
    cv2.imwrite('./test_aligned_deconv.jpg',result)

三、导向反向传播Guided Backpropagation的效果

下左图为上述代码的输入图像，即原图；下右图为为上述代码的输出图像，即原图经导向反向传播后的图。

四、参考链接

1.http://ddrv.cn/a/145213

2.https://blog.csdn.net/qq_41647438/article/details/109504316

3.https://zhuanlan.zhihu.com/p/479485138?utm_id=0

4.https://blog.csdn.net/cdknight_happy/article/details/108792065