8.4 hook函数与CAM可视化

这节课学习hook函数，用hook函数来提取特征图，进行可视化。

然后介绍CAM算法。

 

一、Hook函数的概念

二、Hook函数

三、特征图可视化

四、CAM（class activation map, 类激活图）

 

 

一、Hook函数的概念

hook函数不改变主题，实现额外功能。hook的翻译是钩子。

在pytorch中，采用动态图机制，运算结束之后，一些中间变量会被释放掉。例如特征图、非叶子节点的梯度。但是有时候，我们想要再提取这些中间变量。这时候，我们就可以使用hook函数。

 

二、Hook函数

在pytorch中提供了4种hook函数。

 

其中第一个是针对Tensor的，剩余三个是针对module的。

 

 

1. Tensor.register_hook()

功能是注册一个反向传播的hook函数。因为张量只有在反向传播的时候，非叶子节点的梯度才会消失，所以才有了这个hook函数。

 

这个hook函数只接受一个参数，就是张量的梯度。输出可以是一个张量，也可以是None。如果输出为张量的话，就会覆盖当前张量的梯度；如果输出为None的话，就不会改变梯度。

 

下面我们结合之前的这个计算图，来学习：

计算图如下：

a = x+w

b = w+1

y = a*b 

 

y对a的梯度= b = w+1 = 2

y对b的梯度=a = x+w = 3

y对w的梯度=b*1 + a*1 = a+b = 5. 

我们知道当反向传播结束后，y对a和b的导数就被释放掉。我们使用hook()函数记录下来。

 

例1：保存非叶子节点的梯度信息。

# -*- coding:utf-8 -*-
"""
@file name  : hook_methods.py
@brief      : pytorch的hook函数
"""
import torch
import torch.nn as nn

import sys, os
hello_pytorch_DIR = os.path.abspath(os.path.dirname(__file__)+os.path.sep+".."+os.path.sep+"..")
sys.path.append(hello_pytorch_DIR)

from tools.common_tools import set_seed

set_seed(1)  # 设置随机种子


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()           #用来保存a的梯度

def grad_hook(grad):      #定义钩子函数：传入一个参数grad
    a_grad.append(grad)   #将这个梯度保存一份

handle = a.register_hook(grad_hook)   #注册hook函数：grad_hook()。register_hook()方法接受一个函数。

y.backward()              #反向传播。当执行完之后a的梯度会被释放掉。

# 查看梯度
print("gradient:", w.grad, x.grad, a.grad, b.grad, y.grad)  #a、b、y三个非叶子节点的梯度应该消失了。
print("a_grad[0]: ", a_grad[0])                             #可以看到保存的梯度信息。

handle.remove()

运行结果：

我们在函数主体backward()上挂上了一个额外的方法，grad_hook()，对于张量a，增加了功能，把它的梯度保存了下来。

所以当backward()结束之后，a的梯度也就保存了下来。

 

例2：