### PyTorch Hook 函数传递参数 在 PyTorch 中，`torch.nn.Module` 提供了注册钩子（hook）的功能来监控或修改前向传播和反向传播中的张量数据。对于希望传递额外参数到 hook 函数的情况，可以通过闭包或者部分应用函数的方式实现。 #### 使用闭包方式传递参数 通过定义一个外部作用域内的变量，在内部定义实际的 hook 函数并访问这些外部变量： ```python import torch import torch.nn as nn class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() self.linear = nn.Linear(10, 5) def get_activation(name): activation = {} def hook(model, input, output): activation[name] = output.detach() return hook, activation model = MyModel() # Registering the forward hook with additional parameter 'l1' forward_hook, activations = get_activation('l1') handle = model.linear.register_forward_hook(forward_hook) input_tensor = torch.randn(1, 10) output = model(input_tensor) print(activations['l1']) handle.remove() # Remove the hook after use ``` 此方法利用 Python 的闭包特性，使得 `get_activation` 返回的 hook 函数能够记住创建它时的作用域环境及其局部变量[^1]。 #### 利用 functools.partial 进行部分应用 另一种解决方案是使用标准库中的 `functools.partial` 来固定某些输入参数: ```python from functools import partial def custom_backward_hook(module, grad_input, grad_output, param=None): print(f"Parameter passed to backward hook: {param}") return tuple(gi.clone().detach() for gi in grad_input) backward_hook_with_param = partial(custom_backward_hook, param="example_value") linear_layer = nn.Linear(10, 5).double() _ = linear_layer.double().register_full_backward_hook(backward_hook_with_param) dummy_input = torch.ones((1, 10), dtype=torch.float64, requires_grad=True) loss = (linear_layer(dummy_input)**2).sum() loss.backward() ``` 这里展示了如何借助 `partial` 将特定值绑定至自定义的 backward 钩子作为附加参数。 这两种技术都可以有效地扩展默认 hook 接口的能力，使其支持更多样化的应用场景需求。