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技术