115 lines
4.6 KiB
Python
115 lines
4.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import functools
|
|
from typing import Any
|
|
from torch import nn
|
|
|
|
|
|
class ForwardHook:
|
|
"""
|
|
Base class for forward hooks.
|
|
Hooks are used in the way:
|
|
modified_args, modified_kwargs = hook.pre_forward(module, *args, **kwargs)
|
|
output = module.forward(*modified_args, **modified_kwargs)
|
|
modified_output = hook.post_forward(module, output)
|
|
"""
|
|
|
|
@classmethod
|
|
def name(cls) -> str:
|
|
raise NotImplementedError
|
|
|
|
def on_attach(self, module: nn.Module): # noqa: B027
|
|
"""Called once when the hook is attached to the module."""
|
|
pass
|
|
|
|
def on_detach(self, module: nn.Module): # noqa: B027
|
|
"""
|
|
Called once when the hook is detached from the module.
|
|
Note: this function is not guaranteed to be called if the module is
|
|
deleted before the hook is detached.
|
|
"""
|
|
pass
|
|
|
|
def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple[Any, ...], dict[str, Any]]:
|
|
"""Called before the module's forward method is executed."""
|
|
return args, kwargs
|
|
|
|
def post_forward(self, module: nn.Module, output: Any) -> Any:
|
|
"""Called after the module's forward method is executed."""
|
|
return output
|
|
|
|
|
|
class ModuleHookManager:
|
|
module_hook_attribute = "_hook_manager"
|
|
|
|
def __init__(self, module: nn.Module):
|
|
self.module = module
|
|
self.forward_hooks: dict[str, ForwardHook] = {}
|
|
self.original_forward = module.forward
|
|
|
|
@classmethod
|
|
def get_from(cls, module: nn.Module) -> "ModuleHookManager | None":
|
|
if hasattr(module, cls.module_hook_attribute):
|
|
return getattr(module, cls.module_hook_attribute)
|
|
return None
|
|
|
|
@classmethod
|
|
def get_from_or_default(cls, module: nn.Module) -> "ModuleHookManager":
|
|
if not hasattr(module, cls.module_hook_attribute):
|
|
setattr(module, cls.module_hook_attribute, cls(module))
|
|
|
|
def forward_hook_wrapper(mod: nn.Module, *args, **kwargs):
|
|
manager: ModuleHookManager = getattr(mod, cls.module_hook_attribute)
|
|
for hook in manager.forward_hooks.values():
|
|
args, kwargs = hook.pre_forward(mod, *args, **kwargs)
|
|
output = manager.original_forward(*args, **kwargs)
|
|
for hook in reversed(manager.forward_hooks.values()):
|
|
output = hook.post_forward(mod, output)
|
|
return output
|
|
|
|
module.forward = functools.partial(forward_hook_wrapper, module)
|
|
|
|
return getattr(module, cls.module_hook_attribute)
|
|
|
|
@staticmethod
|
|
def remove_from_manager(module: nn.Module) -> None:
|
|
if hasattr(module, ModuleHookManager.module_hook_attribute):
|
|
manager: ModuleHookManager = getattr(module, ModuleHookManager.module_hook_attribute)
|
|
module.forward = manager.original_forward
|
|
delattr(module, ModuleHookManager.module_hook_attribute)
|
|
|
|
def _check_manager_attached(self) -> None:
|
|
if not hasattr(self.module, self.module_hook_attribute):
|
|
raise ValueError("ModuleHookManager is not attached to the module.")
|
|
if getattr(self.module, self.module_hook_attribute) is not self:
|
|
raise ValueError("ModuleHookManager attached to the module is different.")
|
|
|
|
def append_forward_hook(self, hook: ForwardHook):
|
|
self._check_manager_attached()
|
|
if hook.name() in self.forward_hooks:
|
|
raise ValueError(f"Hook with name {hook.name()} is already registered.")
|
|
# after python 3.7, dicts maintain insertion order
|
|
self.forward_hooks[hook.name()] = hook
|
|
hook.on_attach(self.module)
|
|
|
|
def replace_forward_hook(self, hook_name: str, new_hook: ForwardHook, run_on_attach: bool = True):
|
|
self._check_manager_attached()
|
|
if hook_name not in self.forward_hooks:
|
|
raise ValueError(f"No hook with name {hook_name} found.")
|
|
old_hook = self.forward_hooks[hook_name]
|
|
if run_on_attach:
|
|
old_hook.on_detach(self.module)
|
|
self.forward_hooks[hook_name] = new_hook
|
|
new_hook.on_attach(self.module)
|
|
|
|
def remove_forward_hook(self, hook_name: str, run_detach: bool = True):
|
|
self._check_manager_attached()
|
|
if hook_name not in self.forward_hooks:
|
|
raise ValueError(f"No hook with name {hook_name} found.")
|
|
if run_detach:
|
|
self.forward_hooks[hook_name].on_detach(self.module)
|
|
del self.forward_hooks[hook_name]
|
|
|
|
def get_forward_hook(self, hook_name: str) -> ForwardHook | None:
|
|
return self.forward_hooks.get(hook_name, None)
|