我了解到前向钩子函数的形式为hook_fn(m,x,y)，其中m指的是模型，x指的是输入，y指的是输出。我想为nn.Transformer编写一个前向钩子函数。
然而，变换器层有两个输入，即src和tgt。例如，>>> out = transformer_model(src, tgt)。那么，我如何区分这些输入呢？

回答：

你的钩子会用x和y的tuple类型调用你的回调函数。正如在torch.nn.Module.register_forward_hook的文档页面中描述的那样（虽然它确实没有详细解释x和y的类型）。

输入只包含传递给模块的位置参数。关键字参数不会传递给钩子，只会传递给forward。[…]。

model = nn.Transformer(nhead=16, num_encoder_layers=12)src = torch.rand(10, 32, 512)tgt = torch.rand(20, 32, 512)

定义你的回调函数：

def hook(module, x, y):    print(f'is tuple={isinstance(x, tuple)} - length={len(x)}')          src, tgt = x      print(f'src: {src.shape}')    print(f'tgt: {tgt.shape}')

将钩子附加到你的nn.Module：

>>> model.register_forward_hook(hook)

进行推理：

>>> out = model(src, tgt)is tuple=True - length=2src: torch.Size([10, 32, 512])tgt: torch.Size([20, 32, 512])