pytorch: hook机制，取网络层的输入输出和梯度

前言

本篇记录在pytorch中，通过hook机制将网络某层的输入输出和梯度取出。

pytorch的hook机制

hook是pytorch中一个独特的机制，可以用于将变量的梯度、网络层的输入输出和梯度“钩出来”保存和修改。

使用hook流程：首先定义hook函数对梯度或者输入输出的操作，然后register到变量或者网络层，最后对网络推理或者变量反向传播激活hook函数生效，去除hook。

register_hook用于变量梯度操作

register_hook接收的hook函数只包含变量梯度这一个参数：

import torch

def hook_function(grad):
    grad += 1

a = torch.tensor([1, 1, 1], requires_grad=True, dtype=torch.float32)
b = a.mean()
h = a.register_hook(hook_function)
b.backward()

print(b.grad) # tensor([1.333, 1.333, 1.333])
h.remove()

尽量不要通过register_hook修改变量梯度。如果要取出变量梯度，可以定义全局变量然后把grad赋过去。

register_forward_hook用于网络层输入输出操作