一、register_hook

import torch

from torch.autograd import Variable

grad_list = []

def print_grad(grad):

    grad_list.append(grad)

x = Variable(torch.randn(2, 1), requires_grad=True)

y = x+2

z = torch.mean(torch.pow(y, 2))

lr = 1e-3

y.register_hook(print_grad)

z.backward()

x.data -= lr*x.grad.data

print(grad_list)

`二、register_forward_hook` & `register_backward_hook`

这两个函数的功能类似于variable函数的register_hook，可在module前向传播或反向传播时注册钩子。

每次前向传播执行结束后会执行钩子函数（hook）。前向传播的钩子函数具有如下形式：hook(module, input, output) -> None，而反向传播则具有如下形式：hook(module, grad_input, grad_output) -> Tensor or None。

钩子函数不应修改输入和输出，并且在使用后应及时删除，以避免每次都运行钩子增加运行负载。钩子函数主要用在获取某些中间结果的情景，如中间某一层的输出或某一层的梯度。这些结果本应写在forward函数中，但如果在forward函数中专门加上这些处理，可能会使处理逻辑比较复杂，这时候使用钩子技术就更合适一些。下面考虑一种场景，有一个预训练好的模型，需要提取模型的某一层（不是最后一层）的输出作为特征进行分类，但又不希望修改其原有的模型定义文件，这时就可以利用钩子函数。下面给出实现的伪代码。

测试LeNet网络

import torch as t

import torch.nn as nn

import torch.nn.functional as F

class LeNet(nn.Module):

    def __init__(self):

        super(LeNet,self).__init__()

        self.conv1 = nn.Conv2d(1, 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):

        x = F.max_pool2d(F.relu(self.conv1(x)),(2,2))

        x = F.max_pool2d(F.relu(self.conv2(x)),2)

        x = x.view(x.size()[0], -1)

        x = F.relu(self.fc1(x))

        x = F.relu(self.fc2(x))

        x = self.fc3(x)

        return x

先模拟一下单次的向前传播，

net = LeNet()

img = t.autograd.Variable((t.arange(32*32*1).view(1,1,32,32)))

net(img)

仿照上面示意，进行钩子注册，获取第一卷积层输出结果，

def hook(module, inputdata, output):

    '''把这层的输出拷贝到features中'''

    print(output.data)

handle = net.conv2.register_forward_hook(hook)

net(img)

# 用完hook后删除

handle.remove()

看看hook能识别什么

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()

可见对于目标层，其输入输出都可以获取到，

None

『PyTorch』第十六弹_hook技术

一、register_hook

二、register_forward_hook & register_backward_hook

测试LeNet网络

看看hook能识别什么

`二、register_forward_hook` & `register_backward_hook`