为了更深入地理解神经网络模型,有时候我们需要观察它训练得到的卷积核、特征图或者梯度等信息,这在CNN可视化研究中经常用到。其中,卷积核最易获取,将模型参数保存即可得到;特征图是中间变量,所对应的图像处理完即会被系统清除,否则将严重占用内存;梯度跟特征图类似,除了叶子结点外,其它中间变量的梯度都被会内存释放,因而不能直接获取。最容易想到的获取方法就是改变模型结构,在forward的最后不但返回模型的预测输出,还返回所需要的特征图等信息。

Pytorch的hook编程可以在不改变网络结构的基础上有效获取、改变模型中间变量以及梯度等信息hook可以提取或改变Tensor的梯度,也可以获取nn.Module的输出和梯度(这里不能改变)。因此有3个hook函数用于实现以上功能:

Tensor.register_hook(hook_fn),

nn.Module.register_forward_hook(hook_fn),

nn.Module.register_backward_hook(hook_fn).


1.Tensor.register_hook(hook_fn)

功能:注册一个反向传播hook函数,用于自动记录Tensor的梯度。
PyTorch对中间变量和非叶子节点的梯度运行完后会自动释放,以减缓内存占用。

In [18]: a = torch.Tensor([1,2]).requires_grad_() 
    ...: b = torch.Tensor([3,4]).requires_grad_() 
    ...: d = torch.Tensor([2]).requires_grad_() 
    ...: c = a + b 
    ...: e = c * d 
    ...: o = e.sum()     

In [19]: o.backward()

In [20]: print(a.grad)
tensor([2., 2.])

In [21]: print(b.grad)
tensor([2., 2.])

In [22]: print(c.grad)
None

In [23]: print(d.grad)
tensor([10.])

In [24]: print(e.grad)
None

In [25]: print(o.grad)
None

可以从程序的输出中看到,a,b,d作为叶子节点,经过反向传播后梯度值仍然保留,而其它非叶子节点的梯度已经被自动释放了,要想得到它们的梯度值,就需要使用hook了。

我们首先自定义一个hook_fn函数,用于记录对Tensor梯度的操作,然后用Tensor.register_hook(hook_fn)对要获取梯度的非叶子结点的Tensor进行注册,然后重新反向传播一次():

In [44]: def hook_fn(grad):
    ...:     print(grad)
    ...:

In [45]: e.register_hook(hook_fn)
Out[45]: <torch.utils.hooks.RemovableHandle at 0x1d139cf0a88>

In [46]: o.backward()
tensor([1., 1.])

参考:

PyTorch之HOOK——获取神经网络特征和梯度的有效工具 - 简书

model.apply(fn)或net.apply(fn) - 码农教程

pytorch的hook机制之register_forward_hook - 知乎

register_forward_hook_夜信_的博客-CSDN博客_register_forward_hook 

11-06 11:44