pytorch学习笔记（七）：pytorchhook和关于pytorchbackward过程的理解⽔平有限，如有错误，请不吝指正，谢谢！
pytorch 的 hook 机制
在看pytorch官⽅⽂档的时候，发现在nn.Module部分和Variable部分均有hook的⾝影。感到很神奇，因为在使⽤tensorflow的时候没有碰到过这个词。所以打算⼀探究竟。
Variable 的 hook
register_hook(hook)
注册⼀个backward钩⼦。
每次gradients被计算的时候，这个hook都被调⽤。hook应该拥有以下签名：
hook(grad) -> Variable or None
hook不应该修改它的输⼊，但是它可以返回⼀个替代当前梯度的新梯度。
这个函数返回⼀个 句柄(handle)。它有⼀个⽅法 ve()，可以⽤这个⽅法将hook从module移除。
例⼦：
import torch
v = sor([0,0,0], requires_grad=True, dtype=torch.float32)
h = v.register_hook(lambda grad: grad *2)# double the gradient
v.sor([1,1,1], dtype=torch.float32))
# 先计算原始梯度，再进hook，获得⼀个新梯度。
ad.data)
tensor([2., 2., 2.])
nn.Module的hook
register_forward_hook(hook)
在module上注册⼀个forward hook。
这⾥要注意的是，hook 只能注册到 Module 上，即，仅仅是简单的 op 包装的 Module，⽽不是我们继承 Module时写的那个类，我们继承 Module写的类叫做 Container。
每次调⽤forward()计算输出的时候，这个hook就会被调⽤。它应该拥有以下签名：
hook(module, input, output) -> None
hook不应该修改 input和output的值。 这个函数返回⼀个 句柄(handle)。它有⼀个⽅法 ve()，可以⽤这个⽅法
将hook从module移除。
看这个解释可能有点蒙逼，但是如果要看⼀下nn.Module的源码怎么使⽤hook的话，那就乌云尽散了。
先看 register_forward_hook
def register_forward_hook(self, hook):
handle = hooks.RemovableHandle(self._forward_hooks)
self._forward_hooks[handle.id]= hook
return handle
这个⽅法的作⽤是在此module上注册⼀个hook，函数中第⼀句就没必要在意了，主要看第⼆句，是把注册的hook保存在_forward_hooks字典⾥。
再看 nn.Module 的__call__⽅法（被阉割了，只留下需要关注的部分）：
def__call__(self,*input,**kwargs):
result = self.forward(*input,**kwargs)
for hook in self._forward_hooks.values():
#将注册的hook拿出来⽤
hook_result = hook(self,input, result)
...
return result
可以看到，当我们执⾏model(x)的时候，底层⼲了以下⼏件事：
调⽤ forward ⽅法计算结果
判断有没有注册 forward_hook，有的话，就将 forward 的输⼊及结果作为hook的实参。然后让hook⾃⼰⼲⼀些不可告⼈的事情。看到这，我们就明⽩hook签名的意思了，还有为什么hook不能修改input的output的原因。
⼩例⼦：
import torch
from torch import nn
import torch.functional as F
from torch.autograd import Variable
def for_hook(module,input, output):
print(module)
for val in input:
print("input val:",val)
for out_val in output:
print("output val:", out_val)
class Model(nn.Module):
def__init__(self):
super(Model, self).__init__()
def forward(self, x):
return x+1
model = Model()
x = Variable(torch.FloatTensor([1]), requires_grad=True)
handle = ister_forward_hook(for_hook)
print(model(x))
register_backward_hook
在module上注册⼀个bachward hook。此⽅法⽬前只能⽤在Module上，不能⽤在Container上，当Module的forward函数中只有⼀
个Function的时候，称为Module，如果Module包含其它Module，称之为Container
每次计算module的inputs的梯度的时候，这个hook会被调⽤。hook应该拥有下⾯的signature。
hook(module, grad_input, grad_output) -> Tensor or None
如果module有多个输⼊输出的话，那么grad_input grad_output将会是个tuple。
hook不应该修改它的arguments，但是它可以选择性的返回关于输⼊的梯度，这个返回的梯度在后续的计算中会替代grad_input。
这个函数返回⼀个 句柄(handle)。它有⼀个⽅法 ve()，可以⽤这个⽅法将hook从module移除。
从上边描述来看，backward hook似乎可以帮助我们处理⼀下计算完的梯度。看下⾯nn.Module中register_backward_hook⽅法的实现，和register_forward_hook⽅法的实现⼏乎⼀样，都是⽤字典把注册的hook保存起来。
def register_backward_hook(self, hook):
handle = hooks.RemovableHandle(self._backward_hooks)
self._backward_hooks[handle.id]= hook
return handle
先看个例⼦来看⼀下hook的参数代表了什么：
import torch
from torch.autograd import Variable
import Parameter
as nn
import math
def bh(m,gi,go):
print("Grad Input")
print(gi)
print("Grad Output")
print(go)
return gi[0]*0,gi[1]*0
class Linear(nn.Module):
def__init__(self, in_features, out_features, bias=True):
super(Linear, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = Parameter(torch.Tensor(out_features, in_features)) if bias:
self.bias = Parameter(torch.Tensor(out_features))
else:
def reset_parameters(self):
stdv =1./ math.sqrt(self.weight.size(1))
self.weight.data.uniform_(-stdv, stdv)
if self.bias is not None:
self.bias.data.uniform_(-stdv, stdv)
def forward(self,input):
if self.bias is None:
return self._backend.Linear()(input, self.weight)
else:
return self._backend.Linear()(input, self.weight, self.bias)
x=Variable(torch.FloatTensor([[1,2,3]]),requires_grad=True)
mod=Linear(3,1, bias=False)
out=mod(x)
print(['*']*20)
print("x.grad", x.grad)
print(ad)
Grad Input
(Variable containing:
1.00000e-02 *
5.1902 -2.3778 -4.4071
[torch.FloatTensor of size 1x3]
, Variable containing:
0.1000 0.2000 0.3000
[torch.FloatTensor of size 1x3]
)
Grad Output
(Variable containing:
variable怎么记0.1000
[torch.FloatTensor of size 1x1]
,)
['*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*']
0 -0 -0
[torch.FloatTensor of size 1x3]
Variable containing:
0 0 0
[torch.FloatTensor of size 1x3]
可以看出，grad_in保存的是，此模块Function⽅法的输⼊的值的梯度。grad_out保存的是，此模块forward⽅法返回值的梯度。我们不能在grad_in上直接修改，但是我们可以返回⼀个新的new_grad_in作为Function⽅法inputs的梯度。
上述代码对variable和module同时注册了backward hook，这⾥要注意的是，⽆论是module hook还是variable hook，最终还是注册
到Function上的。这点通过查看Varible的register_hook源码和Module的__call__源码得知。
Module的register_backward_hook的⾏为在未来的⼏个版本可能会改变
BP过程中Function中的动作可能是这样的
class Function:
def__init__(self):
...
def forward(self, inputs):
...
return outputs
def backward(self, grad_outs):
...
return grad_ins
def_backward(self, grad_outs):
hooked_grad_outs = grad_outs
for hook in hook_in_outputs:
hooked_grad_outs = hook(hooked_grad_outs)
grad_ins = self.backward(hooked_grad_outs)
hooked_grad_ins = grad_ins
for hook in hooks_in_module:
hooked_grad_ins = hook(hooked_grad_ins)
return hooked_grad_ins
关于pytorch run_backward()的可能实现猜测为。
def run_backward(variable, gradient):
creator = ator
if creator is None:
return
grad_ins = creator._backward(gradient)
vars= creator.saved_variables
for var, grad in zip(vars, grad_ins):
run_backward(var, ad)
中间Variable的梯度在BP的过程中是保存到GradBuffer中的(C++源码中可以看到), BP完会释放. 如果retain_grads=True的话，就不会被释放
版权声明：本站内容均来自互联网，仅供演示用，请勿用于商业和其他非法用途。如果侵犯了您的权益请与我们联系QQ:729038198，我们将在24小时内删除。
发表评论