pytorch hook使用

因为pytorch会自动舍弃图计算的中间结果，因此想要获取这些数值就须要使用钩子函数。

钩子函数包括Variable的钩子和nn.Module钩子，用法类似。

import torch from torch.autograd import Variable grad_list = [] grad_listx = [] def print_grad(grad): grad_list.append(grad) def print_gradx(grad): grad_listx.append(grad) x = Variable(torch.randn(2, 1), requires_grad=True) y = x*x + 2 z = torch.mean(torch.pow(y, 2)) lr = 1e-3 y.register_hook(print_grad) x.register_hook(print_gradx) z.backward() x.data -= lr * x.grad.data print("x.grad.data-------------") print(x.grad.data) print("y-------------") print(grad_list) print("x-------------") print(grad_listx)

- 输出： 记录了y的梯度，而后x.data=记录x的梯度

/opt/conda/bin/python2.7 /root/rjw/pytorch_test/pytorch_exe03.py x.grad.data-------------

 32.3585
 14.8162 [torch.FloatTensor of size 2x1] y------------- [Variable containing: 7.1379
 4.5970 [torch.FloatTensor of size 2x1] ] x------------- [Variable containing: 32.3585
 14.8162 [torch.FloatTensor of size 2x1] ] Process finished with exit code 0

register_forward_hook & register_backward_hook

• 这两个函数的功能相似于variable函数的register_hook，可在module前向传播或反向传播时注册钩子。每次前向传播执行结束后会执行钩子函数（hook）。前向传播的钩子函数具备以下形式：hook(module, input, output) -> None，而反向传播则具备以下形式：hook(module, grad_input, grad_output) -> Tensor or None。
• 钩子函数不该修改输入和输出，而且在使用后应及时删除，以免每次都运行钩子增长运行负载。钩子函数主要用在获取某些中间结果的情景，如中间某一层的输出或某一层的梯度。这些结果本应写在forward函数中，但若是在forward函数中专门加上这些处理，可能会使处理逻辑比较复杂，这时候使用钩子技术就更合适一些。下面考虑一种场景，有一个预训练好的模型，须要提取模型的某一层（不是最后一层）的输出做为特征进行分类，但又不但愿修改其原有的模型定义文件，这时就能够利用钩子函数。
• 『PyTorch』第十六弹_hook技术