pytorch中autograd以及hook函数详解

原文链接：https://cloud.tencent.com/developer/article/1122582

在这篇文章中：

• 前言
• 正文
  • register hook
  • register_backward_hook
• 后记
• 参考链接：

有些公式为图片，如果这个页面加载不出来，请看这里:https://oldpan.me/archives/pytorch-autograd-hook

前言

pytorch中的Autograd mechanics(自动求梯度机制)是实现前向以及后向反馈运算极为重要的一环，pytorch官方专门针对这个机制进行了一个版块的讲解：

“This note will present an overview of how autograd works and records the operations. It’s not strictly necessary to understand all this, but we recommend getting familiar with it, as it will help you write more efficient, cleaner programs, and can aid you in debugging.”

地址在这里:https://pytorch.org/docs/stable/notes/autograd.html#autograd-mechanics，当然，官方只是说个大概，而且官方还说–不必非要理解它。但是，如果你想要利用神经网络实现一些比较“高级”的功能时，这些概念该理解还是要理解的，在算法中有很多地方会使用到它。

注意，本文所写遵循pytorch的0.4.0版本，在pytorch的0.4.0版本中，Variable和tensor合并，一般在使用中tensor即Variable，也就是不要需要Variable去对tensor进行wrap了。

关于pytorch0.4.0版本的信息请看这里：https://oldpan.me/archives/pytorch-v0-4-0-release。

正文

自动求导求梯度机制相关的一个参数我们应该都熟悉，requires_grad。

当在定义一个tensor的时候并且将requires_grad设置为True。这个tensor就拥有自动求梯度：

>>> x = torch.randn(5, 5) # requires_grad=False by default
>>> y = torch.randn(5, 5) # requires_grad=False by default
>>> z = torch.randn((5, 5), requires_grad=True)
>>> a = x + y
>>> a.requires_grad
False
>>> b = a + z
>>> b.requires_grad
True

这是官方的示例程序，只要有一个tensor的requires_grad设置为True，那么接下来的计算中所有相关的tensor都会支持自动求导求梯度。

关于自动求导求梯度的一些信息请看这里：https://oldpan.me/archives/pytroch-torch-autograd-backward。

register hook

但是自动求导的机制有个我们需要注意的地方：

In[2]: import torch
In[3]: x = torch.tensor([1,2],dtype=torch.float32,requires_grad=True)
In[4]: y = x * 2
In[5]: z = torch.mean(y)
In[6]: z
Out[6]: tensor(3.)
In[7]: z.backward()
In[8]: x.grad
Out[8]: tensor([ 1.,  1.])
In[9]: y.grad # 应该为(0.5,0.5)
In[10]: z.grad # 应该为1

上面的代码中，我们输入x是一个拥有两个元素的向量(x1,x2)，y=(x1*2,x2*2)=(y1,y2) , z = (y1+y2)/2。

最终也就是z = (x1*2 + x2*2)/2，显然x1和x2对输出的导数是(1,1)，但是如果我们看y和z的导数会发现什么也没有输出。

这是为什么，是因为在自动求导机制中只保存叶子节点，也就是中间变量在计算完成梯度后会自动释放以节省空间，所以上面代码我们在计算过程中只得到了z对x的梯度。

但是问题来了，有没有办法可以得到y和z的梯度呢，当然是有的。

这就需要我们的hook函数了：

register hook (hook)[source] 这个函数属于torch.tensor类，这个函数在与这个tensor梯度计算的时候就会执行，这个函数的参数hook是一个函数，这个函数应该是以下的形式：

hook(grad) -> Tensor or None 。grad是这个tensor的梯度，该函数返回grad，我们可以改变这个hook函数的返回值，但是不能改变其参数。

In[2]: import torch
In[3]: x = torch.tensor([1,2],dtype=torch.float32,requires_grad=True)
In[4]: y = x * 2
In[5]: y.requires_grad
Out[5]: True
In[6]: y.register_hook(print)
Out[6]: <torch.utils.hooks.RemovableHandle at 0x7f765e876f60>
In[7]: z = torch.mean(y)
In[8]: z.backward()
tensor([ 0.5000,  0.5000])

在看上面的函数，我们通过对y进行register_hook引入print这个函数，print即是简单的打印，将y相关的grad打印出来，结果我们看到了，在z.backward()执行的时候，y的hook函数也执行了，打印出了y关于输出z的梯度，也就是之前那段代码中的(0.5,0.5)。

register_backward_hook

之前说的是tensor中的register_hook,现在说的这个函数是module类里面的hook函数，module即我们平常使用pytorch定义神经网络层时需要的模板类。register_backward_hook和register_forward_hook是差不多的，一个在backward中执行一个在forward中执行，这里只讲解下register_backward_hook函数。

该函数的形式是register_backward_hook(hook)，同样参数是一个hook函数，hook函数的形式为： hook(module, grad_input, grad_output) -> Tensor or None

register_backward_hook函数同样在module输入的梯度进行计算的时候会执行，注意hook函数中的grad_input和grad_output参数格式不可以改变，但是在hook函数中可以对grad_input参数进行修改并返回一个新的自定义的grad_input，以便在某些算法中实现不同的功能。注意当输入或者输出为多个时，grad_input和grad_output也会变成多个，格式为tuples。

这里演示一个例子：

首先输入 x = [1,2,3,4] 为一个拥有四个元素的变量，这四个元素分别为x1 x2 x3 x4

,然后经过两个运算：

。我们设置w都为8。偏置b为2

。

这样输入一个四维变量的x然后我们得到一维变量的z。下面是通过module的实现代码。

（以下分析需要精力投入）

import torch
import torch.nn as nn

device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)
class MyMul(nn.Module):

def forward(self, input):

out = input  2
return out
class MyMean(nn.Module):            # 自定义除法module

def forward(self, input):

out = input/4
return out

def tensor_hook(grad):
print(‘tensor hook’)
print(‘grad:’, grad)
return grad
class MyNet(nn.Module):

def init(self):
super(MyNet, self).init()

self.f1 = nn.Linear(4, 1, bias=True)

self.f2 = MyMean()

self.weight_init()

def forward(self, input):

self.input = input

output = self.f1(input) # 先进行运算1，后进行运算2

output = self.f2(output)
return output

def weight_init(self):

self.f1.weight.data.fill_(8.0) # 这里设置Linear的权重为8

self.f1.bias.data.fill_(2.0) # 这里设置Linear的bias为2

def my_hook(self, module, grad_input, grad_output):
print(‘doing my_hook’)
print(‘original grad:’, grad_input)
print(‘original outgrad:’, grad_output)

# grad_input = grad_input[0]self.input   # 这里把hook函数内对grad_input的操作进行了注释，

# grad_input = tuple([grad_input])        # 返回的grad_input必须是tuple，所以我们进行了tuple包装。

# print(‘now grad:’, grad_input)
return grad_input
if name == ‘main’:

input = torch.tensor([1, 2, 3, 4], dtype=torch.float32, requires_grad=True).to(device)

net = MyNet()

net.to(device)

net.register_backward_hook(net.my_hook) # 这两个hook函数一定要result = net(input)执行前执行，因为hook函数实在forward的时候进行绑定的

input.register_hook(tensor_hook)

result = net(input)
print(‘result =’, result)

result.backward()
print(‘input.grad:’, input.grad)
for param in net.parameters():
print(’{}:grad->{}’.format(param, param.grad))

上面代码中我们定义了MyNet()，MyNet()中有两个运算module，然后在input和MyNet()中有使用了hook函数，将其grad打印出来。

但是有一点需要注意，在自己设计module的时候，module的运算单元必须也是继承module的，如果你只是简单的这样写：

class MyNet(nn.Module):

def init(self):
…

def forward(self, input):

self.input = input

output = input  2

output = output / 4
return output

虽然在forward中进行了运算，但是这样在backward后这些运算不会参与自动求导，也就是说这些运算不在自动求导求梯度树中，所以必须要继承mudule去定义运算单元。

（上面这段话收回）

好了，执行完上面的代码，输出为：

result = tensor([ 20.5000], device=‘cuda:0’)

doing my_hook

original grad: (tensor([ 0.2500], device=‘cuda:0’),)

original outgrad: (tensor([ 1.], device=‘cuda:0’),)

tensor hook

grad: tensor([ 2., 2., 2., 2.], device=‘cuda:0’)

input.grad: None

Parameter containing:
tensor([[ 8., 8., 8., 8.]], device=‘cuda:0’):grad->tensor([[ 0.2500, 0.5000, 0.7500, 1.0000]], device=‘cuda:0’)

Parameter containing:
tensor([ 2.], device=‘cuda:0’):grad->tensor([ 0.2500], device=‘cuda:0’)

好了我们来按顺序根据输出分析一下整个过程。

结果为result=20.5很容易看出来，(1+2+3+4)8 + 2/4 = 20.5。

然后是MyNet()中的hook函数首先执行，注意输出的original grad是0.25，我们的输入x有4个元素然后输出z为一个标量，但是输出这个grad是0.25只有一个元素。这个0.25是怎么来的，我们之前的两个运算分别是：

。

。

这个0.25是第二个式子

中，z对y的梯度。而original outgrad就是z对z的梯度当然是1了。

接下来是tensor hook，输出的grad: tensor([ 2., 2., 2., 2.], device=’cuda:0′)，这里输出的是z对x的梯度。我们根据前面两个式子很容易得出是[2,2,2,2]。

但我们发现在backward后打印出来input.grad: None。按理说这个应该input的grad应该存在，但是这里打印出来却是None,很有可能是bug问题，这里暂时不讨论。

最后打印的是Linear运算单元的权重和其grad，我们发现w(1-4)权重的grad是[ 0.2500, 0.5000, 0.7500, 1.0000]，这个是怎么得来的，这里是z对w的权重而不是对输入x的，比如w的grad就是

，拿

来举例就是

。

而另一个则是偏置z对偏置b的梯度，这个很容易得到是0.25。

我们再修改一下上面程序中

 output = self.f1(input)

output = self.f2(output)

的顺序，将output = self.f1(input) ==》 output = self.f2(input) 然后将 output = self.f2(input) 》 output = self.f1(input)。

得到的结果为：

result = tensor([ 22.], device=‘cuda:0’)

doing my_hook

original grad: (tensor([ 1.], device=‘cuda:0’), tensor([ 1.], device=‘cuda:0’))

original outgrad: (tensor([ 1.], device=‘cuda:0’),)

tensor hook

grad: tensor([ 2.,  2.,  2.,  2.], device=‘cuda:0’)

input.grad: None

Parameter containing:
tensor([[ 8.,  8.,  8.,  8.]], device=‘cuda:0’):grad->tensor([[ 0.2500,  0.5000,  0.7500,  1.0000]], device=‘cuda:0’)

Parameter containing:
tensor([ 2.], device=‘cuda:0’):grad->tensor([ 1.], device=‘cuda:0’)

发现original grad: (tensor([ 1.], device=’cuda:0′), tensor([ 1.], device=’cuda:0′))发生了变化。

我们再修改一下使forward层为：

 output = self.f2(input)

output = self.f1(output)/4

发现结果变成：

result = tensor([ 5.5000], device=‘cuda:0’)

doing my_hook

original grad: (tensor([ 0.2500], device=‘cuda:0’),)

original outgrad: (tensor([ 1.], device=‘cuda:0’),)

tensor hook

grad: tensor([ 0.5000,  0.5000,  0.5000,  0.5000], device=‘cuda:0’)

input.grad: None

Parameter containing:
tensor([[ 8.,  8.,  8.,  8.]], device=‘cuda:0’):grad->tensor([[ 0.0625,  0.1250,  0.1875,  0.2500]], device=‘cuda:0’)

Parameter containing:
tensor([ 2.], device=‘cuda:0’):grad->tensor([ 0.2500], device=‘cuda:0’)

其他地方不看，还是看original grad: (tensor([ 0.2500], device=’cuda:0′),)，发现从两个tensor变成了一个tensor。

发生上面现象的原因其实是pytorch的一个bug，一个问题，在pytorch的issue中有这样一个回答：

大意是我们的hook函数只会绑定在module中最后一个执行函数上，上面的MyNet在forward函数进行修改后，最后一个执行函数f1或f2发生了变化，所以导致的结果不同：

当–

 output = self.f1(input)

output = self.f2(output)

时，

最后的f2为z = y /4 所以为0.25.

之后我们改成：（f1和f2顺序互换）

 output = self.f2(input)

output = self.f1(output)

那么最后的f2为一个Linear(4,1)。这个Linear有bias也就是这个公式为 z = y1 + y2 + y3 + y4 + b (其中y = wx)。而hook将上面的式子解析成了z = (y1 + y2 + y3 + y4) + b  <> m + b。

也就是两个式子m和b（m为 (y1 + y2 + y3 + y4)集合）所以得到了两个grad ==> z对 m 和 b 的，从式子z = m + b 很容易得到

。也就对应着之前说的 (tensor([ 1.], device=’cuda:0′), tensor([ 1.], device=’cuda:0′))。

最后我们又将forward中的执行函数改成：

 output = self.f2(input)

output = self.f1(output)/4

在f1，第二个式子中除以了4,也就是实 z = y1 + y2 + y3 + y4 + b 变为z = (y1 + y2 + y3 + y4 + b)/4。这样将 (y1 + y2 + y3 + y4 + b)看成n。于是乎，z = n /4。

然后结果显而易见，original grad: (tensor([ 0.2500], device=’cuda:0′),)。

这些是pytorch设计中的一个bug，设计者建议使用tensor的hook而不建议使用module的hook大概是这个原因，但是我们只要多注意一下，知道这些bug就可以不必犯错。

后记

说了这么多，回到之前提到的require_grad参数。在平时设计神经网络的时候并没有特意去设置require_grad这个参数，这是为什么。因为我们平时进行训练的是神经网络的权重，神经网络的权重才是我们要训练的参数，并不是输入也不是其他东西。在pytorch中，在你设计了一个神经网络层后，这个层中的参数默认是可以进行梯度运算的：

# 这里定义一个自编码器的网络层
class Autoencoder(nn.Module):

def init(self):
super(Autoencoder, self).init()

self.encoder = nn.Sequential(
_ConvLayer(3, 128),
_ConvLayer(128, 256),
_ConvLayer(256, 512),
_ConvLayer(512, 1024),
Flatten(),

nn.Linear(1024  4  4,1024),

nn.Linear(1024,1024  4  4),
Reshape(),
_UpScale(1024, 512),
)
…

打印出这个net的encoder部分的参数是否可以requires_grad

for param in model.encoder.parameters():
print(param.requires_grad)

这个结果显然是：

True

True

True

True

True

True

True

True

True

True

True

True

True

True

也就是说，你设计的net里面的权重参数默认都是可以进行自动梯度求导的，我们平常的loss.backward()中反向求导中的所要更新的值也就是net中的权重参数值。