一个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=<AddmmBackward0>)
下面是另一个例子,保存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