『PyTorch』第十六弹_hook技术

由于pytorch会自动舍弃图计算的中间结果，所以想要获取这些数值就需要使用钩子函数。

钩子函数包括Variable的钩子和nn.Module钩子，用法相似。

一、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)

[Variable containing:
 1.5653
 3.5175
[torch.FloatTensor of size 2x1]
]

二、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函数中专门加上这些处理，可能会使处理逻辑比较复杂，这时候使用钩子技术就更合适一些。下面考虑一种场景，有一个预训练好的模型，需要提取模型的某一层（不是最后一层）的输出作为特征进行分类，但又不希望修改其原有的模型定义文件，这时就可以利用钩子函数。下面给出实现的伪代码。

model = VGG()
features = t.Tensor()
def hook(module, input, output):
    '''把这层的输出拷贝到features中'''
    features.copy_(output.data)

handle = model.layer8.register_forward_hook(hook)
_ = model(input)
# 用完hook后删除
handle.remove()

测试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)

Variable containing:

Columns 0 to 7
 27.6373 -13.4590  23.0988 -16.4491  -8.8454 -15.6934  -4.8512   1.3490

Columns 8 to 9
  3.7801 -15.9396
[torch.FloatTensor of size 1x10]

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

def hook(module, inputdata, output):
    '''把这层的输出拷贝到features中'''
    print(output.data)

handle = net.conv2.register_forward_hook(hook)
net(img)
# 用完hook后删除
handle.remove()

……

……

[torch.FloatTensor of size 1x16x10x10]

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

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

Model(
)
input val: Variable containing:
           1
          [torch.FloatTensor of size 1]

output val: Variable containing:
            2
           [torch.FloatTensor of size 1]

Variable containing:
 2
[torch.FloatTensor of size 1]