欢迎关注 “小白玩转Python”，发现更多 “有趣”

使用过深度学习的人都知道，有时候调试模型是非常困难的。张量的不匹配、梯度爆炸，以及其他无数的问题都会让你大吃一惊。解决这些问题需要细微的观察这些模型。最基本的方法包括在forward()方法中添加print语句或引入断点。但是这也相当麻烦，因为需要猜测哪里出现了问题。

现在有了一个解决方案：hooks。这些是特定的函数，可以附加到每一层，并在每次使用层时调用。它们基本上允许冻结特定模块的正向或反向传递的执行，并处理其输入和输出。

下面让我们详细介绍一下！

Hooks速成小课堂

钩子只是一个带有预定义签名的可调用对象，它可以注册到任何nn.Module。当在模块上使用触发器方法(即forward()或backward())时，模块本身及其输入和可能的输出将传递给钩子，在计算进行到下一个模块之前执行。

在 PyTorch 中，可以将钩子注册为：

· forward prehook(在前向传播之前执行)

· forward hook(在前向传播之后执行)

· backward hook(在后向传播之后执行)

乍一看可能很复杂，让我们来看一个具体的例子！

例：保存每个卷积层的输出

假设我们要查看 ResNet34框架中每个卷积层的输出。这项工作非常适合使用hooks。在下一部分中，我将向您展示如何执行这一操作。

我们的模型定义如下：

import torch
from torchvision.models import resnet34
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = resnet34(pretrained=True)
model = model.to(device)

创建一个hook来保存输出非常简单，对于我们的目的来说，一个基本的可调用对象就足够了。

class SaveOutput:def __init__(self):self.outputs = []  def __call__(self, module, module_in, module_out):self.outputs.append(module_out)def clear(self):self.outputs = []

SaveOutput的一个实例将只记录前向传播过程的输出张量并将其存储在一个列表中。

可以使用 register_forward_hook(hook)方法注册前向钩子。(对于其他类型的钩子，我们有 register_backward_hook 和 register_forward_pre_hook。)这些方法的返回值是hook句柄，可用于从模块中删除hook。

现在我们将hook注册到每个卷积层。

save_output = SaveOutput()
hook_handles = []
for layer in model.modules():if isinstance(layer, torch.nn.modules.conv.Conv2d):handle = layer.register_forward_hook(save_output)hook_handles.append(handle)

完成后，钩子将在每个卷积层的每个前向传递后被调用。为了测试它，我们将使用下面的图像。

前向传播：

from PIL import Image
from torchvision import transforms as T
image = Image.open('cat.jpg')
transform = T.Compose([T.Resize((224, 224)), T.ToTensor()])
X = transform(image).unsqueeze(dim=0).to(device)
out = model(X)

正如预期的那样，输出被正确存储。

>>> len(save_output.outputs)
36

通过查看此列表中的张量，我们可以将网络所看到的可视化。

出于好奇，我们可以看看网络接下来会发生什么。如果我们深入网络，学习到的特征会越来越高。例如，有一个过滤器似乎是负责检测眼睛。

超越

当然，这只是冰山一角。Hooks可以做的远不止简单地存储中间层的输出。例如，神经网络剪枝，这是一种策略，以减少参数的数目，也可以通过hooks实现。

总而言之，如果你想增强你的工作，学会使用hooks是一种非常有用的技巧。有了这些，你就能做得更多，做得更有效。

·  END  ·

HAPPY LIFE