pytorch学习笔记(七):pytorch hook 和 关于pytorch backward过程的理解

pytorch 的 hook 机制

在看pytorch官方文档的时候,发现在nn.Module部分和Variable部分均有hook的身影。感到很神奇,因为在使用tensorflow的时候没有碰到过这个词。所以打算一探究竟。

Variable 的 hook

register_hook(hook)

注册一个backward钩子。

每次gradients被计算的时候,这个hook都被调用。hook应该拥有以下签名:

hook(grad) -> Variable or None

hook不应该修改它的输入,但是它可以返回一个替代当前梯度的新梯度。

这个函数返回一个 句柄(handle)。它有一个方法 handle.remove(),可以用这个方法将hookmodule移除。

例子:

v = Variable(torch.Tensor([0, 0, 0]), requires_grad=True)
h = v.register_hook(lambda grad: grad * 2)  # double the gradient
v.backward(torch.Tensor([1, 1, 1]))
#先计算原始梯度,再进hook,获得一个新梯度。
print(v.grad.data)
h.remove()  # removes the hook
 2
 2
 2
[torch.FloatTensor of size 3]

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不应该修改 inputoutput的值。 这个函数返回一个 句柄(handle)。它有一个方法 handle.remove(),可以用这个方法将hookmodule移除。

看这个解释可能有点蒙逼,但是如果要看一下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不能修改inputoutput的原因。

小例子:

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 = model.register_forward_hook(for_hook)
print(model(x))
handle.remove()

register_backward_hook

module上注册一个bachward hook此方法目前只能用在Module上,不能用在Container上,当Module的forward函数中只有一个Function的时候,称为Module,如果Module包含其它Module,称之为Container

每次计算moduleinputs的梯度的时候,这个hook会被调用。hook应该拥有下面的signature

hook(module, grad_input, grad_output) -> Tensor or None

如果module有多个输入输出的话,那么grad_input grad_output将会是个tuplehook不应该修改它的arguments,但是它可以选择性的返回关于输入的梯度,这个返回的梯度在后续的计算中会替代grad_input

这个函数返回一个 句柄(handle)。它有一个方法 handle.remove(),可以用这个方法将hookmodule移除。

从上边描述来看,backward hook似乎可以帮助我们处理一下计算完的梯度。看下面nn.Moduleregister_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
from torch.nn import Parameter
import torch.nn 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:
            self.register_parameter('bias', None)
        self.reset_parameters()

    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)
mod.register_backward_hook(bh) # 在这里给module注册了backward hook

out=mod(x)
out.register_hook(lambda grad: 0.1*grad) #在这里给variable注册了 hook
out.backward()
print(['*']*20)
print("x.grad", x.grad)
print(mod.weight.grad)
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:
 0.1000
[torch.FloatTensor of size 1x1]
,)
['*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*', '*']
x.grad Variable containing:
 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的梯度。

上述代码对variablemodule同时注册了backward hook,这里要注意的是,无论是module hook还是variable hook,最终还是注册到Function上的。这点通过查看Varibleregister_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 = variable.creator
    if creator is None:
        variable.grad = variable.hook(gradient)
        return 
    grad_ins = creator._backward(gradient)
    vars = creator.saved_variables
    for var, grad in zip(vars, grad_ins):
        run_backward(var, var.grad)

中间Variable的梯度在BP的过程中是保存到GradBuffer中的(C++源码中可以看到), BP完会释放. 如果retain_grads=True的话,就不会被释放

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏阿凯的Excel

或关系模糊匹配求均值(虐心升级版)

上期和大家分享了如何使用数组函数实现或关系求均值。 本期和大家分享进一步的应用,或关系模糊匹配求均值。 如果没看上期直接看本期会比较痛苦,来个传送门! 点击我可...

31560
来自专栏苦逼的码农

【算法实战】生成窗口最大值数组

做算法题了,题的难度我们分为“士,尉,校,将”四个等级。这个算法题的模块是篇幅比较小的那种模块。首先是给出一道题的描述,之后我会用我的想法来做这道题,今天算是算...

13920
来自专栏潇涧技术专栏

Python Algorithms - C2 The basics

本节主要介绍了三个内容:算法渐近运行时间的表示方法、六条算法性能评估的经验以及Python中树和图的实现方式。

13020
来自专栏calmound

HDU 1536 S-Nim

师时隔几个月后,在看博弈NIM终于懂了,做出了第一道根据所有堆的异或和的题 题意:在一个Si容器中,输入可以行走的步数,然后在下面T组测试数据中,输入堆的情况 ...

37260
来自专栏ACM算法日常

确定比赛名次(拓扑排序) - HDU 1285

这次先讲理论,因为拓扑排序在日常工作中用的并不多,甚至于很多人可能忘了计算机中存在这样一种排序。我大概的整理一下拓扑排序的定义和应用,以便看了这...

13020
来自专栏章鱼的慢慢技术路

OpenGL中的二维编程——从简单的矩形开始

20640
来自专栏蜉蝣禅修之道

Max-Min Fairness带宽分配算法

28060
来自专栏机器之心

业界 | 谷歌开源DeepLearn.js:可实现硬件加速的机器学习JavaScript库

选自GitHub 机器之心编译 参与:蒋思源、路雪 deeplearn.js 是一个可用于机器智能并加速 WebGL 的开源 JavaScript 库。de...

39780
来自专栏chenjx85的技术专栏

leetcode-201-数字范围按位与

给定范围 [m, n],其中 0 <= m <= n <= 2147483647,返回此范围内所有数字的按位与(包含 m, n 两端点)。

49420
来自专栏CNN

从Tensorflow模型文件中解析并显示网络结构图(pb模型篇)

Tensorflow官方提供的Tensorboard可以可视化神经网络结构图,但是说实话,我几乎从来不用。主要是因为Tensorboard中查看到的图结构太混乱...

3.6K60

扫码关注云+社区

领取腾讯云代金券