load npu accelerating module

Author: MengqingCao

benchmark.py

@@ -24,9 +24,6 @@
 from timm.optim import create_optimizer_v2
 from timm.utils import setup_default_logging, set_jit_fuser, decay_batch_step, check_batch_size_retry, ParseKwargs,\
     reparameterize_model
-from timm.utils.distributed import is_torch_npu_available
-
-has_torch_npu = is_torch_npu_available()
 
 has_apex = False
 try:

inference.py

@@ -21,9 +21,6 @@
 from timm.layers import apply_test_time_pool
 from timm.models import create_model
 from timm.utils import AverageMeter, setup_default_logging, set_jit_fuser, ParseKwargs
-from timm.utils.distributed import is_torch_npu_available
-
-has_torch_npu = is_torch_npu_available()
 
 
 try:

timm/__init__.py

@@ -2,3 +2,41 @@
 from .layers import is_scriptable, is_exportable, set_scriptable, set_exportable
 from .models import create_model, list_models, list_pretrained, is_model, list_modules, model_entrypoint, \
     is_model_pretrained, get_pretrained_cfg, get_pretrained_cfg_value
+import os
+import json
+import importlib
+import sys
+
+# import device specific accelerator module
+device_extension_info = open("./timm/device_extension.json", 'r')
+device_extension_info = json.load(device_extension_info)
+
+os_var = ""
+for device_key in device_extension_info.keys():
+    os_var_modules = device_extension_info[device_key]
+    os_var += device_key + ':'
+
+    for module in os_var_modules:
+        os_var += module + ':'
+    os_var = os_var[:-1]
+    os_var += ','
+
+os.environ["DEVICE_EXT"] = os_var[:-1]
+
+if os.getenv('DEVICE_EXT'):
+    this_module = sys.modules[__name__]
+    backends = os.getenv('DEVICE_EXT').split(',')
+    for backend in backends:
+        module_info = backend.split(':')
+        module_name = module_info[1].strip()
+        module_alias = list()
+        if len(module_info) > 2:
+            for i in range(2, len(module_info)):
+                module_alias.append(module_info[i].strip())
+        try:
+            extra_module = importlib.import_module(module_name)
+            for alia in module_alias:
+                setattr(this_module, alia, extra_module)
+            print(module_alias)
+        except ImportError:
+            pass

timm/device_extension.json

@@ -0,0 +1,3 @@
+{
+    "ascend_npu_modules": ["torch_npu"]
+}

timm/models/_builder.py

@@ -15,7 +15,6 @@
 from timm.models._pretrained import PretrainedCfg
 from timm.models._prune import adapt_model_from_file
 from timm.models._registry import get_pretrained_cfg
-from timm.utils.distributed import is_torch_npu_available
 
 _logger = logging.getLogger(__name__)
 
@@ -24,7 +23,6 @@
 _DOWNLOAD_PROGRESS = False
 _CHECK_HASH = False
 _USE_OLD_CACHE = int(os.environ.get('TIMM_USE_OLD_CACHE', 0)) > 0
-has_torch_npu = is_torch_npu_available()
 
 __all__ = ['set_pretrained_download_progress', 'set_pretrained_check_hash', 'load_custom_pretrained', 'load_pretrained',
            'pretrained_cfg_for_features', 'resolve_pretrained_cfg', 'build_model_with_cfg']

timm/models/_factory.py

@@ -7,9 +7,6 @@
 from ._hub import load_model_config_from_hf
 from ._pretrained import PretrainedCfg
 from ._registry import is_model, model_entrypoint, split_model_name_tag
-from timm.utils.distributed import is_torch_npu_available
-
-has_torch_npu = is_torch_npu_available()
 
 
 __all__ = ['parse_model_name', 'safe_model_name', 'create_model']

timm/utils/__init__.py

@@ -4,7 +4,7 @@
 from .cuda import ApexScaler, NativeScaler
 from .decay_batch import decay_batch_step, check_batch_size_retry
 from .distributed import distribute_bn, reduce_tensor, init_distributed_device,\
-    world_info_from_env, is_distributed_env, is_primary, is_torch_npu_available
+    world_info_from_env, is_distributed_env, is_primary
 from .jit import set_jit_legacy, set_jit_fuser
 from .log import setup_default_logging, FormatterNoInfo
 from .metrics import AverageMeter, accuracy

timm/utils/distributed.py

@@ -2,8 +2,6 @@
 
 Hacked together by / Copyright 2020 Ross Wightman
 """
-import importlib
-import importlib.metadata as imp_meta
 import logging
 import os
 from typing import Optional
@@ -48,19 +46,6 @@ def is_primary(args, local=False):
     return is_local_primary(args) if local else is_global_primary(args)
 
 
-def is_torch_npu_available():
-    _torch_npu_available = importlib.util.find_spec("torch_npu") is not None
-    if _torch_npu_available:
-        try:
-            torch_npu_version = imp_meta.version("torch_npu")
-            import torch_npu  # noqa: F401
-            torch.npu.set_device(0)
-            _logger.info(f"torch_npu version {torch_npu_version} is available.")
-        except ImportError:
-            _torch_npu_available = False
-    return _torch_npu_available
-
-
 def is_distributed_env():
     if 'WORLD_SIZE' in os.environ:
         return int(os.environ['WORLD_SIZE']) > 1
@@ -170,7 +155,11 @@ def init_distributed_device_so(
     if 'cuda' in device:
         assert torch.cuda.is_available(), f'CUDA is not available but {device} was specified.'
     if 'npu' in device:
-        assert is_torch_npu_available(), f'NPU is not available but {device} was specified.'
+        try:
+            TORCH_NPU_AVAILABLE = torch.npu.is_available()
+            assert TORCH_NPU_AVAILABLE, f'NPU is not available but {device} was specified.'
+        except ImportError:
+            _logger.info(f"NPU is not available but {device} was specified.")
 
     if distributed and device != 'cpu':
         device, *device_idx = device.split(':', maxsplit=1)

validate.py

@@ -27,9 +27,6 @@
 from timm.models import create_model, load_checkpoint, is_model, list_models
 from timm.utils import accuracy, AverageMeter, natural_key, setup_default_logging, set_jit_fuser, \
     decay_batch_step, check_batch_size_retry, ParseKwargs, reparameterize_model
-from timm.utils.distributed import is_torch_npu_available
-
-has_torch_npu = is_torch_npu_available()
 
 try:
     from apex import amp
@@ -399,8 +396,12 @@ def _try_run(args, initial_batch_size):
         try:
             if torch.cuda.is_available() and 'cuda' in args.device:
                 torch.cuda.empty_cache()
-            if torch.npu.is_available() and 'npu' in args.device:
-                torch.npu.empty_cache()
+            if 'npu' in args.device:
+                try:
+                    torch.npu.is_available()
+                    torch.npu.empty_cache()
+                except ImportError:
+                    _logger.info("NPU is not available but {args.device} was specified.")
             results = validate(args)
             return results
         except RuntimeError as e: