PyTorch 12. hook的用法

hook

1. 由于pytorch会自动舍弃图计算的中间结果，所以想要这些数值就需要使用钩子函数。钩子函数包括Variable的钩子和nn.Module钩子，用法相似。
2. 在使用hook函数时，不应该修改它的输入，但是它可以返回一个替代当前梯度的新梯度，即，使用该函数会返回一个梯度值

register_hook

针对Tensor变量的hook函数

import torch

grad_list = []
def print_grad(grad):
	grad_list.append(grad)

x = torch.randn(2,1)
y = x+2
y.register_hook(print_grad)
y.backward()

在整个网络进行反向传递时，运行到钩子函数注册的变量后，会保存该变量的梯度，调用print_grad函数，将梯度添加到grad_list中

register_forward_hook

针对网络层的hook函数,具体的可视化博客，可以参考我的另一篇文档，可视化特定层

钩子函数不应该修改输入和输出，并且在使用后应及时删除，以避免每次都运行钩子增加运行负载。钩子函数主要用在获取某些中间结果的情景，如中间某一层的输出或某一层的梯度。

import torch
model = VGG()
features = torch.Tensor()
def hook(module, input, output):
	# 把这一层的输出拷贝到features中
	features.copy_(output.data)

handle = model.layer8.register_forward_hook(hook)
_ = model(input)
# 用完hook后删除
handle.remove()

register_backward_hook

首先介绍Container的概念：当Module的forward函数中只有一个Function的时候，称为Module,如果Module包含其他Module,称之为Container.

在module上注册一个backward hook。此方法目标只能用在Module上，不能用在Container上。

每次计算module的inputs的梯度时，这个hook会被调用

hook(module,grad_input,grad_output)->Tensor or None

求取模块的梯度，与register_grad类似