为了更深入地理解神经网络模型,有时候我们需要观察它训练得到的卷积核、特征图或者梯度等信息,这在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) - 码农教程