pytorch hook学习

register_hook

import torch

x = torch.Tensor([0,1,2,3]).requires_grad_()

y = torch.Tensor([4,5,6,7]).requires_grad_()

w = torch.Tensor([1,2,3,4]).requires_grad_()

z = x+y;

o = w.matmul(z) # o = w(x+y) 中间变量z

o.backward()

print(x.grad,y.grad,z.grad,w.grad,o.grad)

这里的o和z都是中间变量，不是通过指定值来定义的变量，所以是中间变量，所以pytorch并不存储这些变量的梯度。

对于中间变量z，hook的使用方式为: z.register_hook(hook_fn),其中 hook_fn为一个用户自定义的函数，其签名为：hook_fn(grad) -> Tensor or None。

它的输入为变量 z 的梯度，输出为一个 Tensor 或者是 None （None 一般用于直接打印梯度）。反向传播时，梯度传播到变量 z，再继续向前传播之前，将会传入 hook_fn。如果 hook_fn的返回值是 None，那么梯度将不改变，继续向前传播，如果 hook_fn的返回值是 Tensor 类型，则该 Tensor 将取代 z 原有的梯度，向前传播。

import torch

x = torch.Tensor([0,1,2,3]).requires_grad_()

y = torch.Tensor([4,5,6,7]).requires_grad_()

w = torch.Tensor([1,2,3,4]).requires_grad_()

z = x+y;

def hook_fn(grad):

    print(grad)

    return None

z.register_hook(hook_fn)

o = w.matmul(z) # o = w(x+y) 中间变量z

o.backward()

print(x.grad,y.grad,w.grad,z.grad,o.grad)

register_forward_hook

register_forward_hook的作用是获取前向传播过程中，各个网络模块的输入和输出。对于模块 module，其使用方式为：module.register_forward_hook(hook_fn) 。其中 hook_fn的签名为：

hook_fn(module, input, output) -> None

import torch

from torch import nn

class Model(nn.Module):

    def __init__(self):

        super(Model,self).__init__()

        self.fc1 = nn.Linear(3,4) # WT * X + bias

        self.relu1 = nn.ReLU()

        self.fc2 = nn.Linear(4,1)

        self.init()

    def init(self):

        with torch.no_grad():

            # WT * X + bias，所以W为4*3的矩阵，bias为1*4

            self.fc1.weight = torch.nn.Parameter(

                torch.Tensor([[1., 2., 3.],

                              [-4., -5., -6.],

                              [7., 8., 9.],

                              [-10., -11., -12.]]))

            self.fc1.bias = torch.nn.Parameter(torch.Tensor([1.0, 2.0, 3.0, 4.0]))

            self.fc2.weight = torch.nn.Parameter(torch.Tensor([[1.0, 2.0, 3.0, 4.0]]))

            self.fc2.bias = torch.nn.Parameter(torch.Tensor([1.0]))

    def forward(self,x):

        o = self.fc1(x)

        o = self.relu1(o)

        o = self.fc2(o)

        return o

def hook_fn_forward(module,input,output):

    print(module)

    print(input)

    print(output)

model = Model()

modules = model.named_children()

'''

named_children()

Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

'''

for name,module in modules:

    # 这里的name就是自己定义的self.xx的xx。如上面的fc1，fc2.

    # module代指的就是fc1代表的module等等

    module.register_forward_hook(hook_fn_forward)

x = torch.Tensor([[1.0,1.0,1.0]]).requires_grad_()

o = model(x)

o.backward()

 '''

 Linear(in_features=3, out_features=4, bias=True)

(tensor([[1., 1., 1.]], requires_grad=True),)

tensor([[  7., -13.,  27., -29.]], grad_fn=<AddmmBackward>)

ReLU()

(tensor([[  7., -13.,  27., -29.]], grad_fn=<AddmmBackward>),)

tensor([[ 7.,  0., 27.,  0.]], grad_fn=<ReluBackward0>)

Linear(in_features=4, out_features=1, bias=True)

(tensor([[ 7.,  0., 27.,  0.]], grad_fn=<ReluBackward0>),)

tensor([[89.]], grad_fn=<AddmmBackward>)

 '''

register_backward_hook

理同前者。得到梯度值。

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

上面的代码forward全部替换为backward，结果为：

'''

Linear(in_features=4, out_features=1, bias=True)

(tensor([1.]), tensor([[1., 2., 3., 4.]]), tensor([[ 7.],

        [ 0.],

        [27.],

        [ 0.]]))

(tensor([[1.]]),)

ReLU()

(tensor([[1., 0., 3., 0.]]),)

(tensor([[1., 2., 3., 4.]]),)

Linear(in_features=3, out_features=4, bias=True)

(tensor([1., 0., 3., 0.]), tensor([[22., 26., 30.]]), tensor([[1., 0., 3., 0.],

        [1., 0., 3., 0.],

        [1., 0., 3., 0.]]))

(tensor([[1., 0., 3., 0.]]),)

'''

register_backward_hook只能操作简单模块，而不能操作包含多个子模块的复杂模块。如果对复杂模块用了 backward hook，那么我们只能得到该模块最后一次简单操作的梯度信息。

可以这么用,可以得到一个模块的梯度。

class Mymodel(nn.Module):

	......

model = Mymodel()

model.register_backward_hook(hook_fn_backward)

None

[torch] pytorch hook学习

pytorch hook学习

register_hook

register_forward_hook

register_backward_hook