tensorplay.utils.hooks

Classes

class BackwardHook [source]

BackwardHook(module, user_hooks, user_pre_hooks)

A wrapper class to implement nn.Module backward hooks.

It handles:

• Ignoring non-Tensor inputs and replacing them by None before calling the user hook
• Generating the proper Node to capture a set of Tensor's gradients
• Linking the gradients captures for the outputs with the gradients captured for the input
• Calling the user hook once both output and input gradients are available

Methods

__init__(self, module, user_hooks, user_pre_hooks) [source]

Initialize self. See help(type(self)) for accurate signature.

---

setup_input_hook(self, args) [source]

---

setup_output_hook(self, args) [source]

---

class RemovableHandle [source]

RemovableHandle(hooks_dict: Any, *, extra_dict: Any = None) -> None

A handle which provides the capability to remove a hook.

Args

• hooks_dict (dict): A dictionary of hooks, indexed by hook id.
• extra_dict (Union[dict, List[dict]]): An additional dictionary or list of dictionaries whose keys will be deleted when the same keys are removed from hooks_dict.

Methods

__init__(self, hooks_dict: Any, *, extra_dict: Any = None) -> None [source]

Initialize self. See help(type(self)) for accurate signature.

---

remove(self) -> None [source]

---

Functions

unserializable_hook() [source]

unserializable_hook(f)

Mark a function as an unserializable hook with this decorator.

This suppresses warnings that would otherwise arise if you attempt to serialize a tensor that has a hook.

warn_if_has_hooks() [source]

warn_if_has_hooks(tensor)