Lines Matching full:hook
83 def __init__(self, hook: Callable, module: Optional["Module"] = None):
84 self.hook: Callable = hook
85 functools.update_wrapper(self, hook)
97 raise RuntimeError("You are trying to call the hook of a dead Module!")
98 return self.hook(module, *args, **kwargs)
99 return self.hook(*args, **kwargs)
102 result = {"hook": self.hook, "with_module": self.with_module}
109 self.hook = state["hook"]
115 "You are trying to revive the hook of a dead Module!"
134 hook: Callable[..., None],
136 r"""Register a buffer registration hook common to all modules.
142 The hook will be called every time :func:`register_buffer` is invoked.
145 hook(module, name, buffer) -> None or new buffer
147 The hook can modify the input or return a single modified value in the hook.
151 a handle that can be used to remove the added hook by calling
155 _global_buffer_registration_hooks[handle.id] = hook
160 hook: Callable[..., None],
162 r"""Register a module registration hook common to all modules.
168 The hook will be called every time :func:`register_module` is invoked.
171 hook(module, name, submodule) -> None or new submodule
173 The hook can modify the input or return a single modified value in the hook.
177 a handle that can be used to remove the added hook by calling
181 _global_module_registration_hooks[handle.id] = hook
186 hook: Callable[..., None],
188 r"""Register a parameter registration hook common to all modules.
194 The hook will be called every time :func:`register_parameter` is invoked.
197 hook(module, name, param) -> None or new parameter
199 The hook can modify the input or return a single modified value in the hook.
203 a handle that can be used to remove the added hook by calling
207 _global_parameter_registration_hooks[handle.id] = hook
211 def register_module_forward_pre_hook(hook: Callable[..., None]) -> RemovableHandle:
212 r"""Register a forward pre-hook common to all modules.
219 The hook will be called every time before :func:`forward` is invoked.
222 hook(module, input) -> None or modified input
226 The hook can modify the input. User can either return a tuple or a
227 single modified value in the hook. We will wrap the value into a tuple
230 This hook has precedence over the specific module hooks registered with
235 a handle that can be used to remove the added hook by calling
239 _global_forward_pre_hooks[handle.id] = hook
244 hook: Callable[..., None],
248 r"""Register a global forward hook for all the modules.
255 The hook will be called every time after :func:`forward` has computed an output.
258 hook(module, input, output) -> None or modified output
262 The hook can modify the output. It can modify the input inplace but
267 hook (Callable): The user defined hook to be registered.
268 always_call (bool): If ``True`` the ``hook`` will be run regardless of
273 a handle that can be used to remove the added hook by calling
276 This hook will be executed before specific module hooks registered with
282 _global_forward_hooks[handle.id] = hook
289 hook: Callable[["Module", _grad_t, _grad_t], Union[None, _grad_t]],
291 r"""Register a backward hook common to all the modules.
299 a handle that can be used to remove the added hook by calling
307 "global Module hook. Please use only one of them."
313 _global_backward_hooks[handle.id] = hook
318 hook: Callable[["Module", _grad_t], Union[None, _grad_t]],
320 r"""Register a backward pre-hook common to all the modules.
335 a handle that can be used to remove the added hook by calling
340 _global_backward_pre_hooks[handle.id] = hook
345 hook: Callable[["Module", _grad_t, _grad_t], Union[None, _grad_t]],
347 r"""Register a backward hook common to all the modules.
362 a handle that can be used to remove the added hook by calling
370 "global Module hook. Please use only one of them."
376 _global_backward_hooks[handle.id] = hook
571 for hook in _global_buffer_registration_hooks.values():
572 output = hook(self, name, tensor)
625 for hook in _global_parameter_registration_hooks.values():
626 output = hook(self, name, param)
653 for hook in _global_module_registration_hooks.values():
654 output = hook(self, name, module)
1344 hook: Callable[["Module", _grad_t], Union[None, _grad_t]],
1347 r"""Register a backward pre-hook on the module.
1349 The hook will be called every time the gradients for the module are computed.
1350 The hook should have the following signature::
1352 hook(module, grad_output) -> tuple[Tensor] or None
1354 The :attr:`grad_output` is a tuple. The hook should
1360 For technical reasons, when this hook is applied to a Module, its forward function will
1369 hook (Callable): The user-defined hook to be registered.
1370 prepend (bool): If true, the provided ``hook`` will be fired before
1373 ``hook`` will be fired after all existing ``backward_pre`` hooks
1381 a handle that can be used to remove the added hook by calling
1386 self._backward_pre_hooks[handle.id] = hook
1392 self, hook: Callable[["Module", _grad_t, _grad_t], Union[None, _grad_t]]
1394 r"""Register a backward hook on the module.
1401 a handle that can be used to remove the added hook by calling
1414 self._backward_hooks[handle.id] = hook
1419 hook: Callable[["Module", _grad_t, _grad_t], Union[None, _grad_t]],
1422 r"""Register a backward hook on the module.
1424 The hook will be called every time the gradients with respect to a module
1425 are computed, i.e. the hook will execute if and only if the gradients with
1426 respect to module outputs are computed. The hook should have the following
1429 hook(module, grad_input, grad_output) -> tuple(Tensor) or None
1432 with respect to the inputs and outputs respectively. The hook should
1440 For technical reasons, when this hook is applied to a Module, its forward function will
1449 hook (Callable): The user-defined hook to be registered.
1450 prepend (bool): If true, the provided ``hook`` will be fired before
1453 ``hook`` will be fired after all existing ``backward`` hooks on
1461 a handle that can be used to remove the added hook by calling
1474 self._backward_hooks[handle.id] = hook
1515 "in future versions. This hook will be missing some of the grad_output. "
1532 "in future versions. This hook will be missing some of the grad_input. "
1547 "Using a non-full backward hook when outputs are nested in python data structure "
1548 "is deprecated and will be removed in future versions. This hook will be missing "
1555 … "Using a non-full backward hook when outputs are generated by different autograd Nodes "
1556 "is deprecated and will be removed in future versions. This hook will be missing "
1562 # At this point the grad_output part of the hook will most likely be correct
1569 … "Using a non-full backward hook when the forward contains multiple autograd Nodes "
1570 … "is deprecated and will be removed in future versions. This hook will be missing "
1579 hook: Union[
1590 r"""Register a forward pre-hook on the module.
1592 The hook will be called every time before :func:`forward` is invoked.
1597 passed to the hooks and only to the ``forward``. The hook can modify the
1599 hook. We will wrap the value into a tuple if a single value is returned
1600 (unless that value is already a tuple). The hook should have the
1603 hook(module, args) -> None or modified input
1605 If ``with_kwargs`` is true, the forward pre-hook will be passed the
1606 kwargs given to the forward function. And if the hook modifies the
1607 input, both the args and kwargs should be returned. The hook should have
1610 hook(module, args, kwargs) -> None or a tuple of modified input and kwargs
1613 hook (Callable): The user defined hook to be registered.
1614 prepend (bool): If true, the provided ``hook`` will be fired before
1617 ``hook`` will be fired after all existing ``forward_pre`` hooks
1623 with_kwargs (bool): If true, the ``hook`` will be passed the kwargs
1629 a handle that can be used to remove the added hook by calling
1635 self._forward_pre_hooks[handle.id] = hook
1645 hook: Union[
1654 r"""Register a forward hook on the module.
1656 The hook will be called every time after :func:`forward` has computed an output.
1660 passed to the hooks and only to the ``forward``. The hook can modify the
1662 forward since this is called after :func:`forward` is called. The hook
1665 hook(module, args, output) -> None or modified output
1667 If ``with_kwargs`` is ``True``, the forward hook will be passed the
1669 output possibly modified. The hook should have the following signature::
1671 hook(module, args, kwargs, output) -> None or modified output
1674 hook (Callable): The user defined hook to be registered.
1675 prepend (bool): If ``True``, the provided ``hook`` will be fired
1678 ``hook`` will be fired after all existing ``forward`` hooks on
1684 with_kwargs (bool): If ``True``, the ``hook`` will be passed the
1687 always_call (bool): If ``True`` the ``hook`` will be run regardless of
1693 a handle that can be used to remove the added hook by calling
1703 self._forward_hooks[handle.id] = hook
1764 for hook_id, hook in (
1769 args_kwargs_result = hook(self, args, kwargs) # type: ignore[misc]
1775 "forward pre-hook must return None or a tuple "
1779 args_result = hook(self, args)
1792 for hook_id, hook in (
1796 # mark that always called hook is run
1801 hook_result = hook(self, args, kwargs, result)
1803 hook_result = hook(self, args, result)
1825 for hook in non_full_backward_hooks:
1826 grad_fn.register_hook(_WrappedHook(hook, self))
1849 for hook_id, hook in _global_forward_hooks.items():
1852 hook_result = hook(self, args, result) # type: ignore[possibly-undefined]
1856 … warnings.warn("global module forward hook with ``always_call=True`` raised an exception "
1860 for hook_id, hook in self._forward_hooks.items():
1864 … hook_result = hook(self, args, kwargs, result) # type: ignore[possibly-undefined]
1866 … hook_result = hook(self, args, result) # type: ignore[possibly-undefined]
1870 … warnings.warn("module forward hook with ``always_call=True`` raised an exception "
1977 for hook in _global_module_registration_hooks.values():
1978 output = hook(self, name, value)
1988 for hook in _global_module_registration_hooks.values():
1989 output = hook(self, name, value)
2045 def _register_state_dict_hook(self, hook): argument
2046 r"""Register a post-hook for the :meth:`~torch.nn.Module.state_dict` method.
2049 hook(module, state_dict, prefix, local_metadata) -> None or state_dict
2055 if getattr(hook, "_from_public_api", False):
2057 "Cannot register the same function as the state dict post hook that was "
2061 self._state_dict_hooks[handle.id] = hook
2064 def register_state_dict_post_hook(self, hook): argument
2065 r"""Register a post-hook for the :meth:`~torch.nn.Module.state_dict` method.
2068 hook(module, state_dict, prefix, local_metadata) -> None
2076 # the state_dict by the hook. However, since hooks registered via both these
2078 # cannot be changed due to many dependencies on it, we mark a hook
2082 # for the root module, otherwise, we ensure that the hook returns None.
2083 hook._from_public_api = True
2085 self._state_dict_hooks[handle.id] = hook
2088 def register_state_dict_pre_hook(self, hook): argument
2089 r"""Register a pre-hook for the :meth:`~torch.nn.Module.state_dict` method.
2092 hook(module, prefix, keep_vars) -> None
2098 self._state_dict_pre_hooks[handle.id] = hook
2214 for hook in self._state_dict_pre_hooks.values():
2215 hook(self, prefix, keep_vars)
2224 for hook in self._state_dict_hooks.values():
2225 hook_result = hook(self, destination, prefix, local_metadata)
2226 if not getattr(hook, "_from_public_api", False):
2231 raise RuntimeError("state_dict post-hook must return None")
2234 def _register_load_state_dict_pre_hook(self, hook, with_module=False): argument
2238 hook will not take the ``module`` as the first argument whereas
2243 hook (Callable): Callable hook that will be invoked before
2246 instance to the hook as the first parameter.
2250 hook, self if with_module else None
2254 def register_load_state_dict_pre_hook(self, hook): argument
2255 … r"""Register a pre-hook to be run before module's :meth:`~nn.Module.load_state_dict` is called.
2258 …hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs…
2261 hook (Callable): Callable hook that will be invoked before
2264 return self._register_load_state_dict_pre_hook(hook, with_module=True)
2266 def register_load_state_dict_post_hook(self, hook): argument
2267 … r"""Register a post-hook to be run after module's :meth:`~nn.Module.load_state_dict` is called.
2270 hook(module, incompatible_keys) -> None
2272 The ``module`` argument is the current module that this hook is registered
2281 ``strict=True`` are affected by modifications the hook makes to
2288 a handle that can be used to remove the added hook by calling
2292 self._load_state_dict_post_hooks[handle.id] = hook
2340 for hook in self._load_state_dict_pre_hooks.values():
2341 hook(
2376 # non-lazy modules, since they dont have the hook to do the checks
2554 # Note that the hook can modify missing_keys and unexpected_keys.
2556 for hook in module._load_state_dict_post_hooks.values():
2557 out = hook(module, incompatible_keys)