Python PyTorch Module.register_full_backward_hook用法及代码示例

本文简要介绍python语言中 torch.nn.Module.register_full_backward_hook 的用法。

用法:

register_full_backward_hook(hook)

返回：

一个句柄，可用于通过调用 handle.remove() 删除添加的钩子

返回类型：

torch.utils.hooks.RemovableHandle

在模块上注册一个后向钩子。

每次计算与模块输入相关的梯度时，都会调用该钩子。挂钩应具有以下签名：

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

grad_input 和 grad_output 是元组，分别包含关于输入和输出的梯度。钩子不应该修改它的参数，但它可以选择返回一个关于输入的新梯度，该梯度将在后续计算中用于代替grad_input。 grad_input 将仅对应于作为位置参数给出的输入，并且所有 kwarg 参数都被忽略。对于所有非张量参数，grad_input 和 grad_output 中的条目将为 None。

由于技术原因，当这个钩子被应用到一个 Module 时，它的 forward 函数将接收传递给 Module 的每个 Tensor 的视图。同样，调用者将收到模块 forward 函数返回的每个张量的视图。

警告

使用后向钩子时不允许就地修改输入或输出，这会引发错误。