pytorch中register_hook以及register_forward

何为叶子节点和非叶子节点

在理解register_hook之前，首先得搞懂什么叶子节点和非叶子节。简单来说叶子节点是有梯度且独立得张量，例如a = torch.tensor(2.0,requires_grad=True)，b= torch.tensor(3.0,requires_grad=True)，非叶子节点是依赖其他张量而得到得张量如c = a+b。
判断是叶子节点还是非叶子节点可以使用 is_leaf来判断一个张量是叶子节点还是非叶子节点。

import torch
a = torch.tensor(2.0,requires_grad=True)
b = torch.tensor(3.0,requires_grad=True)
print(a.is_leaf)
print(b.is_leaf)
c = a +b
print(c.is_leaf)>>> True
>>> True
>>> False

中间张量 c 作为非叶子节点是没有梯度信息得。pytorch默认在梯度反向传播过程中不会记录中间变量梯度信息。而且叶子节点的梯度信息在反向传播流过程中是不允许我们修改的。只能通过print(a.grad)查看张量的梯度信息。
那么，如果我们想查看中间变量 c 以及想改变叶子节点反向传播过程中的梯度值，应该怎么办呢。这时候就要使用register_hook这个钩子函数了。通过一下两段代码看一下钩子函数的主要作用。

register_hook

a = torch.tensor(2.0,requires_grad=True)
b = torch.tensor(3.0,requires_grad=True)
print(a.grad)
print(b.grad)
c = a*b
print(c.grad)  # 由于c是叶子节点，所以他是不记录梯度信息得。前后打印梯度信息都为Noned = torch.tensor(4.0,requires_grad=True)
e = c * d
e.backward()
print(a.grad)
print(b.grad)
print(c.grad)>>>输出
None
None
None
tensor(12.)
tensor(8.)
None

通过上面代码可以看出，c作为中间变量在反向传播过程中不记录梯度信息。c=a*b其中a的梯度就为b的值，b的梯度就是a的值。接下来对中间变量c 使用register_hook，这个函数传入的参数得是一个函数。

import torcha = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(3.0, requires_grad=True)c = a * bdef c_hook(grad):print("c_hook",grad)return grad + 2    # 什么也不返回的话用的是和之前一样的梯度，不对其进行变化。# 在c中，钩子按照有序字典的方式存储，按照存储的前后一次调用
c.register_hook(c_hook)
c.register_hook(lambda grad: print("hello my grad is",grad))
c.retain_grad()   # 存储中间变量的梯度print(a.grad)
print(b.grad)
print(c.grad)c.backward()print(a.grad)
print(b.grad)
print(c.grad)>>>
None
None
None
c_hook tensor(1.)
hello my grad is tensor(3.)
tensor(9.)
tensor(6.)
tensor(3.)

为什么输出会是这样的结果呢，一个张量可以注册多个钩子函数，反向传播过程中按照注册的顺序依次运行。 c.register_hook(c_hook) c.register_hook(lambda grad:)
，这两个函数可以重写c的梯度，第一个函数传入的参数是c的梯度，自身对自身的梯度pytorch中默认为1。所以此时c_hook中传入的grad=1,这个函数返回值为grad+2=3,此时会重写中间变量c的梯度信息。第二个钩子函数传入的函数为匿名函数，这个匿名函数对c的梯度没有进行重写，使用的还是上一个钩子函数重写的值，此使打印信息就为3。最后通过c.retain_grad(）记c的梯度信息。通过这个例子，我稍微懂了点register_hook这个钩子函数的作用，是不是本来不可修改的梯度信息值，通过这个函数修改了呢。

通过一下这个例子比较再来看一下registe_hook函数的作用。

import torch
a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(3.0, requires_grad=True)c = a * bdef c_hook(grad):print("c_hook",grad)return grad + 2    # 什么也不返回的话用的是和之前一样的梯度，不对其进行变化。# 在c中，钩子按照有序字典的方式存储，按照存储的前后一次调用
c.register_hook(c_hook)
c.register_hook(lambda grad: print("hello my grad is",grad))
c.retain_grad()   # 存储中间变量的梯度d = torch.tensor(4.0, requires_grad=True)
d.register_hook(lambda grad: grad + 100)  # 将使用100+grad代替本来返回得梯度值e = c * dprint(a.grad)
print(b.grad)
print(c.grad)
print(d.grad)
print(e.grad)# e.retain_grad()
e.register_hook(lambda grad: grad * 2)
e.retain_grad()e.backward()print(a.grad)
print(b.grad)
print(c.grad)
print(d.grad)
print(e.grad)>>>输出
None
None
None
None
None
c_hook tensor(8.)
hello my grad is tensor(10.)
tensor(30.)
tensor(20.)
tensor(10.)
tensor(112.)
tensor(2.)

这段代码前部分和前面的代码保持一致，后面添加了e = c * d，在反向传播前，毋庸置疑a,b,,c,d,e的梯度都为None。反向传播过程中首先看 e，自身对自身的倒数默认为1，但是e注册的钩子将对原本的梯度 * 2 ，来替代原先的梯度信息，所以打印出的e的梯度信息为2。相应的，e 对 c的梯度信息相应的就变为 2d=8,e对d的梯度信息就变为 2c=12,案例说此使d的梯度信息为12，为什么是112呢，可以看出d注册了一个钩子函数，这个钩子给d原本的梯度信息加了100，来代替旧的梯度信息，所以d的梯度信息为112。由于c注册的钩子函数给他加了2，所以c的梯度信息为10。相应的a b 的梯度就都要乘以c 的梯度信息了。同样，原本不变的梯度信息值在这里都根据register_hook这个函数相应的被重写。

以上就是我根据视频链接对register_hook的理解。

register_forward_hook

register_forward_hook register_forward_pre_hook这个函数主要使用在nn.Module网络中。
第一个函数看名称是用在网络forward之前，第二个是运行在forward之后，举例：

import torch
import torch.nn as nnclass SumNet(nn.Module):def __init__(self):super(SumNet, self).__init__()@staticmethoddef forward(a, b, c):d = a + b + cprint('forward():')print('    a:', a)print('    b:', b)print('    c:', c)print()print('    d:', d)print()return ddef forward_pre_hook(module, input_positional_args):a, b, c = input_positional_argsnew_input_positional_args = a + 10, b,c+10print('forward_pre_hook():')print('    module:', module)print('    input_positional_args:', input_positional_args)print()print('    new_input_positional_args:', new_input_positional_args)print()return new_input_positional_argsdef forward_hook(module, input_positional_args, output):new_output = output + 100print('forward_hook():')print('    module:', module)print('    input_positional_args:', input_positional_args)print('    output:', output)print()print('    new_output:', new_output)print()return new_outputdef main():sum_net = SumNet()sum_net.register_forward_pre_hook(forward_pre_hook)sum_net.register_forward_hook(forward_hook)a = torch.tensor(1.0, requires_grad=True)b = torch.tensor(2.0, requires_grad=True)c = torch.tensor(3.0, requires_grad=True)print('start')print()print('a:', a)print('b:', b)print('c:', c)print()print('before model')print()d = sum_net(a, b, c)   # 前向传播得时候钩子函数起作用了，先是forward_pre_hook,接下来是forward,接下来是forward_hook函数。print('after model')print()print('d:', d)if __name__ == '__main__':main()

输出信息：

starta: tensor(1., requires_grad=True)
b: tensor(2., requires_grad=True)
c: tensor(3., requires_grad=True)before modelforward_pre_hook():module: SumNet()input_positional_args: (tensor(1., requires_grad=True), tensor(2., requires_grad=True), tensor(3., requires_grad=True))new_input_positional_args: (tensor(11., grad_fn=<AddBackward0>), tensor(2., requires_grad=True), tensor(13., grad_fn=<AddBackward0>))forward():a: tensor(11., grad_fn=<AddBackward0>)b: tensor(2., requires_grad=True)c: tensor(13., grad_fn=<AddBackward0>)d: tensor(26., grad_fn=<AddBackward0>)forward_hook():module: SumNet()input_positional_args: (tensor(11., grad_fn=<AddBackward0>), tensor(2., requires_grad=True), tensor(13., grad_fn=<AddBackward0>))output: tensor(26., grad_fn=<AddBackward0>)new_output: tensor(126., grad_fn=<AddBackward0>)after modeld: tensor(126., grad_fn=<AddBackward0>)

分析以上为什么会输出这样的结果，前面提到register_forward_hook这个函数会在网络前向传播前运行，需要两个参数modul 和 input案例中输入为 tensor 1 2 3,经过这个函数给2 3 分别加了10，并且返回了一组新的值，这组值是要传入forward中，可以看出，forward函数打印的a b c 为传入的这组新值，而不是刚开始定义的1 2 3，forward函数运行过程中返回每层的输出会运行forward_hook函数。这个函数主要需要三个参数，module input output。
以下从Lenet网络来使用这个函数：

import torch
import torch.nn as nn
import torch.nn.functional as Fclass LeNet(nn.Module):def __init__(self):super(LeNet, self).__init__()self.conv1 = nn.Conv2d(3, 6, 5)self.conv2 = nn.Conv2d(6, 16, 5)self.fc1 = nn.Linear(16*5*5, 120)self.fc2 = nn.Linear(120, 84)self.fc3 = nn.Linear(84, 10)def forward(self, x):out = self.conv1(x)out = F.relu(out)     out = F.max_pool2d(out, 2)      out = self.conv2(out)out = F.relu(out)  out = F.max_pool2d(out, 2)out = out.view(out.size(0), -1)out = F.relu(self.fc1(out))out = F.relu(self.fc2(out))out = self.fc3(out)return out
model = LeNet()# 分别对model的第一个卷积层和最后一层使用了钩子函数，这样既可以取出对应层的输出。
def hook(model,input_,output):print("最后一层输出：",output.shape)def conv_hook(model,input_,output):print("conv1后",input_[0].shape,output.shape)model.register_forward_hook(hook)
model.conv1.register_forward_hook(conv_hook)img = torch.randn([1,3,32,32])
out_put = model(img)>>>
conv1后 torch.Size([1, 3, 32, 32]) torch.Size([1, 6, 28, 28])
最后一层输出： torch.Size([1, 10])

基于上可以看出给不同层使用钩子函数，可以提取出每一层的输出，并进行相应的处理。

以上就是pytorch中register_hook和register_forward_hook的基本理解。
如果有问题烦请指出加以改正。