pytorch hook使用,『,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技术
    原文作者:pytorch
    原文地址: https://www.cnblogs.com/ranjiewen/p/9221702.html
    本文转自网络文章,转载此文章仅为分享知识,如有侵权,请联系博主进行删除。
点赞