Pytorch中register_forward_hook的用法

一个hook是一个可调用对象，可以通过不同的方式注册到任意nn.Module对象上，当module的触发方法比如（forward()或backward()）被调用后，module本身以及其输入和输出被传递给hook并触发调用。

比如当一个hook通过torch.nn.Module.register_forward_hook()的方式注册到一个module上时，每次其forward函数计算出其输出之后紧接着调用hook。

hook的参数形式是固定的 hook(module, input, output) -> None or modified output，参数名是任意的，传入的参数分别是module本身对象、该module forward前的输入、该module forward后的输出。

hook常用于捕获中间过程，如下是一个简单的例子

import torch
from torch import nn


def forward_hook_fn(module, input, output):
    print('weight', module.weight.data)
    print('bias', module.bias.data)
    print('input', input)
    print('output', output)


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc = nn.Linear(3, 1)
        self.fc.register_forward_hook(forward_hook_fn)
        if hasattr(self.fc, 'weight') and self.fc.weight is not None:
            nn.init.constant_(self.fc.weight, 1)
        if hasattr(self.fc, 'bias') and self.fc.bias is not None:
            nn.init.constant_(self.fc.bias, 0)

    def forward(self, x):
        o = self.fc(x)
        return o


if __name__ == '__main__':
    model = Model()
    x = torch.Tensor([[0.0, 1.0, 2.0]])
    y = model(x)

# 输出
weight: tensor([[1., 1., 1.]])
bias: tensor([0.])
input: (tensor([[0., 1., 2.]]),)
output: tensor([[3.]], grad_fn=)

下面是另一个例子，保存resnet34一次前向过程中所有卷积层的输出

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)


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 = []


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)
print(len(save_output.outputs))

# 输出
# 36

参考

How hooks can improve your workflow significantly 

Module — PyTorch 1.11.0 documentation

MMCV 核心组件分析(六): Hook - 知乎