Pytorch中的hook

由于pytorch的中间变量在完成计算后就会被释放，如果想要获取中间的信息，就需要使用hook相关函数。

首先搭建一个深度学习网络

class Net(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 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):
        out = F.relu(self.conv1(x))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv2(out))
        out = F.max_pool2d(out, 2)
        out = out.view(out.size(0), -1)
        out = F.relu(self.fc1(out))
        out = F.relu(self.fc2(out))
        out = self.fc3(out)
        return out

指定网络中的conv2层添加hook，hook单独写一个函数来实现相关功能
register_hook，对象为Variable对象

• register_forward_hook，对应前向传播，用来导出指定子模块（可以是层、模块等nn.Module类型）的输入输出张量，但只可修改输出，常用来导出或修改卷积特征图

res = []
def hook(module, input, output):
    res.append(output.clone().detach())
    
net = Net()
x = torch.randn(2, 3, 32, 32)
handle = net.conv2.register_forward_hook(hook)
out = net(x)
print(out[0].size())
handle.remove()

• register_backward_hook，对应反向传播，用来导出指定子模块输入输出张量的梯度，但只可修改输入张量的梯度，输出张量梯度不可修改

'''
grad_input:表示输入的梯度，不能修改，但是可以通过return替代，元组包含(bias的梯度，输入x的梯度，权重weight的梯度)
grad_output：元组包含输出y的梯度
如果是没有参数的Module，如ReLU，grad_inpu只包含输入x的梯度
'''
def new_hook(module, grad_input, grad_output):
    print("grad_input: ", grad_input)
    print("grad_output: ", grad_output)
    return grad_input[0] * 0, grad_input[1] * 0, grad_input[2] * 0

x = torch.tensor([[1., 2., 10.]], requires_grad=True)
module = nn.Linear(3, 1)
handle = module.register_backward_hook(new_hook)
y = module(x)
y.backward(retain_graph=True)
print()
handle.remove

通过remove函数取消hook