Optim-wip: Improve ModuleOutputsHook, testing coverage

captum/optim/_core/output_hook.py

@@ -1,5 +1,6 @@
 import warnings
-from typing import Callable, Iterable, Tuple
+from collections import OrderedDict
+from typing import Callable, Dict, Iterable, Optional, Tuple
 from warnings import warn
 
 import torch
@@ -15,6 +16,9 @@ def __init__(self, target_modules: Iterable[nn.Module]) -> None:
 
             target_modules (Iterable of nn.Module): A list of nn.Module targets.
         """
+        for module in target_modules:
+            # Clean up any old hooks that weren't properly deleted
+            _remove_all_forward_hooks(module, "module_outputs_forward_hook")
         self.outputs: ModuleOutputMapping = dict.fromkeys(target_modules, None)
         self.hooks = [
             module.register_forward_hook(self._forward_hook())
@@ -33,13 +37,13 @@ def is_ready(self) -> bool:
 
     def _forward_hook(self) -> Callable:
         """
-        Return the forward_hook function.
+        Return the module_outputs_forward_hook forward hook function.
 
         Returns:
-            forward_hook (Callable): The forward_hook function.
+            forward_hook (Callable): The module_outputs_forward_hook function.
         """
 
-        def forward_hook(
+        def module_outputs_forward_hook(
             module: nn.Module, input: Tuple[torch.Tensor], output: torch.Tensor
         ) -> None:
             assert module in self.outputs.keys()
@@ -57,7 +61,7 @@ def forward_hook(
                     "that you are passing model layers in your losses."
                 )
 
-        return forward_hook
+        return module_outputs_forward_hook
 
     def consume_outputs(self) -> ModuleOutputMapping:
         """
@@ -130,3 +134,53 @@ def __call__(self, input_t: TupleOfTensorsOrTensorType) -> ModuleOutputMapping:
         finally:
             self.layers.remove_hooks()
         return activations_dict
+
+
+def _remove_all_forward_hooks(
+    module: torch.nn.Module, hook_fn_name: Optional[str] = None
+) -> None:
+    """
+    This function removes all forward hooks in the specified module, without requiring
+    any hook handles. This lets us clean up & remove any hooks that weren't property
+    deleted.
+
+    Warning: Various PyTorch modules and systems make use of hooks, and thus extreme
+    caution should be exercised when removing all hooks. Users are recommended to give
+    their hook function a unique name that can be used to safely identify and remove
+    the target forward hooks.
+
+    Args:
+
+        module (nn.Module): The module instance to remove forward hooks from.
+        hook_fn_name (str, optional): Optionally only remove specific forward hooks
+            based on their function's __name__ attribute.
+            Default: None
+    """
+
+    if hook_fn_name is None:
+        warn("Removing all active hooks will break some PyTorch modules & systems.")
+
+    def _remove_hooks(m: torch.nn.Module, name: Optional[str] = None) -> None:
+        if hasattr(module, "_forward_hooks"):
+            if m._forward_hooks != OrderedDict():
+                if name is not None:
+                    dict_items = list(m._forward_hooks.items())
+                    m._forward_hooks = OrderedDict(
+                        [(i, fn) for i, fn in dict_items if fn.__name__ != name]
+                    )
+                else:
+                    m._forward_hooks: Dict[int, Callable] = OrderedDict()
+
+    def _remove_child_hooks(
+        target_module: torch.nn.Module, hook_name: Optional[str] = None
+    ) -> None:
+        for name, child in target_module._modules.items():
+            if child is not None:
+                _remove_hooks(child, hook_name)
+                _remove_child_hooks(child, hook_name)
+
+    # Remove hooks from target submodules
+    _remove_child_hooks(module, hook_fn_name)
+
+    # Remove hooks from the target module
+    _remove_hooks(module, hook_fn_name)