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用于网络层输入输出操作

register_forward_hook接收的hook函数包含网络层，层输入和层输出三个参数：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.linear1 = nn.Linear(3, 10)
        self.linear2 = nn.Linear(10, 1)
    def forward(self, x):
        return F.relu(self.linear2(F.relu(self.linear1(x))))

def hook_forward_function(module, module_input, module_output):
    print(module_input)
    print(module_output)

mlp = MLP()
h1 = mlp.linear2.register_forward_hook(hook_forward_function)
a = torch.tensor([1, 1, 1], requires_grad=True, dtype=torch.float32)
y = mlp(a)

h1.remove()

网络推理到hook注册层时，hook函数会被调用。

register_full_backward_hook用于网络层输入输出梯度操作

与forward hook类似，register_full_backward_hook也接收三参数hook函数，但分别是网络层，层输入梯度和层输出梯度。

import torch
import torch.nn as nn
import torch.nn.functional as F

class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.linear1 = nn.Linear(3, 10)
        self.linear2 = nn.Linear(10, 1)
    def forward(self, x):
        return F.relu(self.linear2(F.relu(self.linear1(x))))

def hook_backward_function(module, module_input_grad, module_output_gard):
    print(module_input_grad)
    print(module_output_grad)

mlp = MLP()
h1 = mlp.register_full_backward_hook(hook_backward_function)
a = torch.tensor([1, 1, 1], requires_grad=True, dtype=torch.float32)
y = mlp(a)
y.backward()

h1.remove()

网络输出backward时，hook函数被调用。