Npu dev

mmcv/device/npu/data_parallel.py

@@ -5,8 +5,8 @@
 
 import torch
 
+from mmcv.device.scatter_gather import scatter_kwargs
 from mmcv.parallel import MMDataParallel
-from .scatter_gather import scatter_kwargs
 
 
 def _check_balance(*args, **kwargs):
@@ -19,7 +19,7 @@ def _check_balance(*args, **kwargs):
 # _check_balance function in DataParallel to make initialization pass.
 for m in sys.modules:
     if m.startswith('torch') or 'mmcv' in m:
-        if getattr(sys.modules[m], '_check_balance', None) is not None:
+        if hasattr(sys.modules[m], '_check_balance'):
             setattr(sys.modules[m], '_check_balance', _check_balance)
 
 

mmcv/device/npu/distributed.py

@@ -1,8 +1,8 @@
 # Copyright Huawei Technologies Co., Ltd. All rights reserved.
 # Copyright (c) OpenMMLab. All rights reserved.
 
+from mmcv.device.scatter_gather import scatter_kwargs
 from mmcv.parallel import MMDistributedDataParallel
-from .scatter_gather import scatter_kwargs
 
 
 class NPUDistributedDataParallel(MMDistributedDataParallel):

mmcv/runner/fp16_utils.py

@@ -10,15 +10,18 @@
 import torch.nn as nn
 from torch.nn.parameter import Parameter
 
-from mmcv.utils import TORCH_VERSION, digit_version
+from mmcv.utils import IS_NPU_AVAILABLE, TORCH_VERSION, digit_version
 from .dist_utils import allreduce_grads as _allreduce_grads
 
 try:
     # If PyTorch version >= 1.6.0, torch.cuda.amp.autocast would be imported
     # and used; otherwise, auto fp16 will adopt mmcv's implementation.
     # Note that when PyTorch >= 1.6.0, we still cast tensor types to fp16
     # manually, so the behavior may not be consistent with real amp.
-    from torch.cuda.amp import autocast
+    if IS_NPU_AVAILABLE:
+        from torch.npu.amp import autocast
+    else:
+        from torch.cuda.amp import autocast
 except ImportError:
     pass