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