【深度之眼】【Pytorch打开第12天】:hook函数与CAM可视化

本文深入探讨了Pytorch中的hook函数机制,包括Tensor.register_hook、Module.register_forward_hook、Module.register_forward_pre_hook和Module.register_backward_hook。文章通过实例展示了如何使用hook函数实现AlexNet第一个卷积层特征图的可视化,并对比了inplace=True和inplace=False的影响。此外,文章还介绍了CAM和Grad-CAM两种类激活图方法,解释了它们的工作原理以及在模型分析中的应用。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

任务:hook函数与CAM(class activation map, 类激活图)

任务简介:学习pytorch的hook函数机制以及CAM可视化算法

详细说明:深入学习了解pytorch的hook函数运行机制,介绍pytorch中提供的4种hook函数

  1. torch.Tensor.register_hook(hook)
  2. torch.nn.Module.register_forward_hook
  3. torch.nn.Module.register_forward_pre_hook
  4. torch.nn.Module.register_backward_hook
    最后,介绍CAM可视化及其改进算法Grad-CAM

实现任务

采用torch.nn.Module.register_forward_hook机制实现AlexNet第一个卷积层输出特征图的可视化,并将/torchvision/models/alexnet.py中第28行改为:nn.ReLU(inplace=False),观察inplace=True与inplace=False的差异。


知识点

Hook函数

Hook函数机制:不改变主体,实现额外功能,像一个挂件,挂钩,hook

  1. torch.Tensor.register_hook(hook)
  2. torch.nn.Module.register_forward_hook
  3. torch.nn.Module.register_forward_pre_hook
  4. torch.nn.Module.register_backward_hook

Tensor.register_hook

  • 功能:注册一个反向传播hook函数
  • Hook函数仅一个输入参数,为张量的梯度

    使用方法:.register_hook(hook_fn),其中hook_fn为一个用户自定义的函数,输出为一个 Tensor 或者是 None (None 一般用于直接打印梯度)。反向传播时,梯度传播到变量 z,再继续向前传播之前,将会传入hook_fn。如果hook_fn的返回值是 None,那么梯度将不改变,继续向前传播,如果hook_fn的返回值是Tensor类型,则该Tensor将取代 z 原有的梯度,向前传播。

# ----------------------------------- 1 tensor hook 1
flag = 0
flag = 1
if flag:
    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)
    a = torch.add(w, x)
    b = torch.add(w, 1)
    y = torch.mul(a, b)

    a_grad = list()

    def grad_hook(grad):
        a_grad.append(grad)

    handle = a.register_hook(grad_hook)

    y.backward()

    print("gradient:", w.grad, x.grad, a.grad, b.grad, y.grad)
    print("a_grad[0]:", a_grad[0])

    handle.remove()
    

输出:

gradient: tensor([5.]) tensor([2.]) None None None
a_grad[0]: tensor([2.])

# ----------------------------------- 2 tensor hook 2
flag = 0
flag = 1
if flag:
    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)
    a = torch.add(w, x)
    b = torch.add(w, 1)
    y = torch.mul(a, b)

    a_grad = list()


    def grad_hook(grad):
        grad *= 2
        return grad * 3


    handle = w.register_hook(grad_hook)

    y.backward()

    print("w.grad:", w.grad)
    handle.remove()
    

输出:w.grad: tensor([30.])

Module.register_forward_hook

  • 功能:注册module的前向传播hook函数
  • 参数:
    • module: 当前网络层
    • input:当前网络层输入数据
    • output:当前网络层输出数据

Module.register_forward_pre_hook

  • 功能:注册module的前向传播前的hook函数
  • 参数:
    • module: 当前网络层
    • input:当前网络层输入数据

Module.register_backward_hook

  • 功能:注册module的反向传播的hook函数
  • 参数:
    • module: 当前网络层
    • grad_input:当前网络层输入梯度数据
    • grad_output:当前网络层输出梯度数据

小实验
在这里插入图片描述

# --------------------------- 3 Module.register_forward_hook and pre hook
# flag = 0
flag = 1
if flag:
    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.conv1 = nn.Conv2d(1, 2
以下是基于pytorch写的hook提取模型特定层特征图并可视化CAM的代码: ```python import torch import torch.nn as nn import cv2 import numpy as np class CAM(): def __init__(self, model, target_layer): self.model = model self.target_layer = target_layer self.features = [] self.grads = [] self.hook_fn = None def hook(self, module, input, output): self.features.append(output.cpu().data.numpy()) def hook_backward(self, module, grad_input, grad_output): self.grads.append(grad_output[0].cpu().data.numpy()) def get_cam(self, input_image, class_idx=None): self.features = [] self.grads = [] self.hook_fn = self.target_layer.register_forward_hook(self.hook) hook_fn_backward = self.target_layer.register_backward_hook(self.hook_backward) input_image = input_image.to(device) self.model.zero_grad() output = self.model(input_image) if class_idx is None: class_idx = np.argmax(output.cpu().data.numpy()) one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32) one_hot[0][class_idx] = 1 one_hot = torch.from_numpy(one_hot).requires_grad_(True) one_hot = torch.sum(one_hot * output) self.model.zero_grad() one_hot.backward() grads_val = self.grads[-1] target = self.features[-1] weights = np.mean(grads_val, axis=(2, 3))[0, :] cam = np.zeros(target.shape[2:], dtype=np.float32) for i, w in enumerate(weights): cam += w * target[0, i, :, :] cam = np.maximum(cam, 0) cam = cv2.resize(cam, input_image.shape[2:]) cam = cam - np.min(cam) cam = cam / np.max(cam) self.hook_fn.remove() hook_fn_backward.remove() return cam # 使用示例: class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) self.fc1 = nn.Linear(128 * 8 * 8, 256) self.fc2 = nn.Linear(256, 10) def forward(self, x): x = nn.functional.relu(self.conv1(x)) x = nn.functional.max_pool2d(x, 2) x = nn.functional.relu(self.conv2(x)) x = nn.functional.max_pool2d(x, 2) x = nn.functional.relu(self.conv3(x)) x = nn.functional.max_pool2d(x, 2) x = x.view(-1, 128 * 8 * 8) x = nn.functional.relu(self.fc1(x)) x = self.fc2(x) return x model = Net() target_layer = model.conv3 cam = CAM(model, target_layer) # 加载图像 image_path = 'test.jpg' image = cv2.imread(image_path) image = cv2.resize(image, (32, 32)) image = np.transpose(image, (2, 0, 1)) image = image.astype(np.float32) / 255. image = torch.from_numpy(image) image = image.unsqueeze(0) # 生成CAM cam_map = cam.get_cam(image) # 可视化CAM heatmap = cv2.applyColorMap(np.uint8(255 * cam_map), cv2.COLORMAP_JET) heatmap = np.float32(heatmap) / 255 cam_image = heatmap + np.float32(image[0]) cam_image = cam_image / np.max(cam_image) cv2.imshow('CAM', cam_image) cv2.waitKey(0) ``` 在上面的代码中,我们定义了一个CAM类,用来提取模型特定层的特征图,并生成对应的CAM图像。CAM类中包含了一个hook函数,用来提取目标层的特征图,以及一个hook_backward函数,用来提取特征图对应的梯度。在get_cam函数中,我们首先将输入图像经过模型前向传播,然后根据输出结果确定目标类别。接着,我们通过反向传播计算目标类别对应的特征图梯度,并利用这个梯度生成CAM图像。最后,我们将CAM图像和原始图像叠加起来,生成可视化CAM图像。
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值