feature: named_hook

easytorch/utils/__init__.py

@@ -2,9 +2,11 @@
 from .timer import Timer, TimePredictor
 from .dist import get_rank, get_local_rank, get_world_size, is_rank, is_master, master_only, get_dist_backend
 from .logging import get_logger
+from .named_hook import NamedForwardHook, NamedBackwardHook
+
 
 __all__ = [
     'set_gpus', 'Timer', 'TimePredictor', 'set_tf32_mode', 'get_logger',
     'get_rank', 'get_local_rank', 'get_world_size', 'is_rank', 'is_master', 'master_only',
-    'get_dist_backend', 'setup_random_seed'
+    'get_dist_backend', 'setup_random_seed', 'NamedForwardHook', 'NamedBackwardHook'
 ]

easytorch/utils/named_hook.py

@@ -0,0 +1,26 @@
+from abc import ABCMeta, abstractmethod
+from typing import Tuple
+
+import torch
+from torch import nn
+
+
+class NamedHook(metaclass=ABCMeta):
+    def __init__(self, name: str):
+        self.name = name
+
+    @abstractmethod
+    def __call__(self, module: nn.Module, *args, **kwargs):
+        pass
+
+
+class NamedForwardHook(NamedHook, metaclass=ABCMeta):
+    @abstractmethod
+    def __call__(self, module: nn.Module, inputs: Tuple[torch.Tensor], outputs: Tuple[torch.Tensor]):
+        pass
+
+
+class NamedBackwardHook(NamedHook, metaclass=ABCMeta):
+    @abstractmethod
+    def __call__(self, module: nn.Module, input_grads: Tuple[torch.Tensor], output_grads: Tuple[torch.Tensor]):
+        pass