[Feature]Add a NPU Hook

Author: xuuyangg

mmengine/device/utils.py

@@ -1,5 +1,4 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import os
 from typing import Optional
 
 import torch
@@ -8,10 +7,6 @@
     import torch_npu  # noqa: F401
     import torch_npu.npu.utils as npu_utils
 
-    # Enable operator support for dynamic shape and
-    # binary operator support on the NPU.
-    npu_jit_compile = bool(os.getenv('NPUJITCompile', False))
-    torch.npu.set_compile_mode(jit_compile=npu_jit_compile)
     IS_NPU_AVAILABLE = hasattr(torch, 'npu') and torch.npu.is_available()
 except Exception:
     IS_NPU_AVAILABLE = False

mmengine/hooks/__init__.py

@@ -7,6 +7,7 @@
 from .iter_timer_hook import IterTimerHook
 from .logger_hook import LoggerHook
 from .naive_visualization_hook import NaiveVisualizationHook
+from .npu_hook import NPUHook
 from .param_scheduler_hook import ParamSchedulerHook
 from .profiler_hook import NPUProfilerHook, ProfilerHook
 from .runtime_info_hook import RuntimeInfoHook
@@ -18,5 +19,5 @@
     'Hook', 'IterTimerHook', 'DistSamplerSeedHook', 'ParamSchedulerHook',
     'SyncBuffersHook', 'EmptyCacheHook', 'CheckpointHook', 'LoggerHook',
     'NaiveVisualizationHook', 'EMAHook', 'RuntimeInfoHook', 'ProfilerHook',
-    'PrepareTTAHook', 'NPUProfilerHook', 'EarlyStoppingHook'
+    'PrepareTTAHook', 'NPUProfilerHook', 'EarlyStoppingHook', 'NPUHook'
 ]

mmengine/hooks/npu_hook.py

@@ -0,0 +1,30 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import os
+
+import torch
+
+from mmengine.registry import HOOKS
+from mmengine.utils import try_import
+from .hook import Hook
+
+torch_npu = try_import('torch_npu')
+
+
+@HOOKS.register_module()
+class NPUHook(Hook):
+    """A hook specific for NPU device."""
+
+    priority = 'HIGHEST'
+
+    def __init__(self) -> None:
+        super().__init__()
+        if torch_npu is None:
+            raise ImportError(f'For availability of {self.__class__.__name__},'
+                              'pleasse install ascend pytorch first.')
+
+    def before_run(self, runner) -> None:
+
+        # Enable operator support for dynamic shape and
+        # binary operator support on the NPU.
+        npu_jit_compile = bool(os.getenv('NPUJITCompile', False))
+        torch.npu.set_compile_mode(jit_compile=npu_jit_compile)

mmengine/runner/runner.py

@@ -423,6 +423,10 @@ def __init__(
         self._hooks: List[Hook] = []
         # register hooks to `self._hooks`
         self.register_hooks(default_hooks, custom_hooks)
+
+        if get_device() == 'npu':
+            self.register_hook({'type': 'NPUHook'})
+
         # log hooks information
         self.logger.info(f'Hooks will be executed in the following '
                          f'order:\n{self.get_hooks_info()}')

mmengine/utils/__init__.py

@@ -7,7 +7,7 @@
                    is_method_overridden, is_seq_of, is_str, is_tuple_of,
                    iter_cast, list_cast, requires_executable, requires_package,
                    slice_list, to_1tuple, to_2tuple, to_3tuple, to_4tuple,
-                   to_ntuple, tuple_cast)
+                   to_ntuple, try_import, tuple_cast)
 from .package_utils import (call_command, get_installed_path, install_package,
                             is_installed)
 from .path import (check_file_exist, fopen, is_abs, is_filepath,
@@ -29,5 +29,5 @@
     'get_git_hash', 'ManagerMeta', 'ManagerMixin', 'Timer', 'check_time',
     'TimerError', 'ProgressBar', 'track_iter_progress',
     'track_parallel_progress', 'track_progress', 'deprecated_function',
-    'apply_to', 'get_object_from_string'
+    'apply_to', 'get_object_from_string', 'try_import'
 ]

mmengine/utils/misc.py

@@ -11,6 +11,7 @@
 from importlib import import_module
 from inspect import getfullargspec, ismodule
 from itertools import repeat
+from types import ModuleType
 from typing import Any, Callable, Optional, Type, Union
 
 
@@ -540,3 +541,20 @@ def get_object_from_string(obj_name: str):
             return obj_cls
         except AttributeError:
             return None
+
+
+def try_import(name: str) -> Optional[ModuleType]:
+    """Try to import a module.
+
+    Args:
+        name (str): Specifies what module to import in absolute or relative
+            terms (e.g. either pkg.mod or ..mod).
+
+    Returns:
+        ModuleType or None: If importing successfully, returns the imported
+        module, otherwise returns None.
+    """
+    try:
+        return import_module(name)
+    except ImportError:
+        return None