pytorch hook 钩子函数的用法

钩子编程（hooking），也称作“挂钩”，是计算机程序设计术语，指通过拦截软件模块间的函数调用、消息传递、事件传递来修改或扩展操作系统、应用程序或其他软件组件的行为的各种技术。处理被拦截的函数调用、事件、消息的代码，被称为钩子（hook）。

Hook 是 PyTorch 中一个十分有用的特性。利用它，我们可以不必改变网络输入输出的结构，方便地获取、改变网络中间层变量的值和梯度。这个功能被广泛用于可视化神经网络中间层的 feature、gradient，从而诊断神经网络中可能出现的问题，分析网络有效性。

本文主要用 hook 函数输出网络执行过程中 forward 和 backward 的执行顺序，以此找到了bug所在。

用法如下：

# 设置hook func
def hook_func(name, module):
    def hook_function(module, inputs, outputs):
        # 请依据使用场景自定义函数
        print(name+' inputs', inputs)
        print(name+' outputs', outputs)
    return hook_function

# 注册正反向hook
for name, module in model.named_modules():
    module.register_forward_hook(hook_func('[forward]: '+name, module))
    module.register_backward_hook(hook_func('[backward]: '+name, module))

如一个简单的 MNIST 手写数字识别的模型结构如下：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

打印模型：

Net(
  (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
  (dropout1): Dropout(p=0.25, inplace=False)
  (dropout2): Dropout(p=0.5, inplace=False)
  (fc1): Linear(in_features=9216, out_features=128, bias=True)
  (fc2): Linear(in_features=128, out_features=10, bias=True)
)

构建hook函数：

# 设置hook func
def hook_func(name, module):
    def hook_function(module, inputs, outputs):
        with open("log_model.txt", 'a+') as f:
            # 请依据使用场景自定义函数
            f.write(name + '   len(inputs): ' + str(len(inputs)) + '\n')
            f.write(name + '   len(outputs):  ' + str(len(outputs)) + '\n')
    return hook_function

# 注册正反向hook
for name, module in model.named_modules():
    module.register_forward_hook(hook_func('[forward]: '+name, module))
    module.register_backward_hook(hook_func('[backward]: '+name, module))

输出的前向和反向传播过程：

[forward]: conv1   len(inputs): 1
[forward]: conv1   len(outputs):  8
[forward]: conv2   len(inputs): 1
[forward]: conv2   len(outputs):  8
[forward]: dropout1   len(inputs): 1
[forward]: dropout1   len(outputs):  8
[forward]: fc1   len(inputs): 1
[forward]: fc1   len(outputs):  8
[forward]: dropout2   len(inputs): 1
[forward]: dropout2   len(outputs):  8
[forward]: fc2   len(inputs): 1
[forward]: fc2   len(outputs):  8
[forward]:    len(inputs): 1
[forward]:    len(outputs):  8
[backward]:    len(inputs): 2
[backward]:    len(outputs):  1
[backward]: fc2   len(inputs): 3
[backward]: fc2   len(outputs):  1
[backward]: dropout2   len(inputs): 1
[backward]: dropout2   len(outputs):  1
[backward]: fc1   len(inputs): 3
[backward]: fc1   len(outputs):  1
[backward]: dropout1   len(inputs): 1
[backward]: dropout1   len(outputs):  1
[backward]: conv2   len(inputs): 2
[backward]: conv2   len(outputs):  1
[backward]: conv1   len(inputs): 2
[backward]: conv1   len(outputs):  1

因为只要模型处于train状态，hook_func 就会执行，导致不断输出 [forward] 和 [backward]，所以将输出内容建议写到文件中，而不是 print

到此这篇关于pytorch hook 钩子函数的用法的文章就介绍到这了,更多相关pytorch hook 钩子函数内容请搜索我们以前的文章或继续浏览下面的相关文章希望大家以后多多支持我们！