PyTorch学习—16.PyTorch中hook函数

一、Hook函数概念

  Hook函数机制：不改变主体，实现额外功能，像一个挂件一样将功能挂到函数主体上。Hook函数与PyTorch中的动态图运算机制有关，因为在动态图计算，在运算结束后，中间变量是会被释放掉的，例如：非叶子节点的梯度。但是，我们往往想要提取这些中间变量，这时，我们就可以采用Hook函数在前向传播与反向传播主体上挂上一些额外的功能（函数），通过这些函数获取中间的梯度，甚至是改变中间的梯度。PyTorch一共提供了四种Hook函数：

• torch.Tensor.register_hook(hook)
• torch.nn.Module.register_forward_hook
• torch.nn.Module.register_forward_pre_hook
• torch.nn.Module.register_backward_hook

一种是针对Tensor，其余三种是针对网络的

二、四种Hook函数介绍

1. Tensor.register_hook

def register_hook(self, hook):
	"""
	接受一个hook函数
	"""
	...

功能：注册一个反向传播hook函数，这是因为张量在反向传播的时候，如果不是叶子节点，它的梯度就会消失。由于反向传播过程中存在数据的释放，所以就有了反向传播的hook函数

• Hook函数仅一个输入参数，为张量的梯度

  hook(grad) -> Tensor or None
  

  下面，我们通过计算图流程来观察张量梯度的获取以及熟悉Hook函数。
$$y=(x+w)\ast (w+1)$$

import torch
import torch.nn as nn

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):
    """
    定义一个hook函数，将梯度存储到列表中
    :param grad: 梯度
    :return:
    """
    a_grad.append(grad)

# 注册一个反向传播的hook函数，功能是将梯度存储到a_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()

tensor([5