init npu

mmcls/apis/train.py

@@ -128,10 +128,14 @@ def train_model(model,
         find_unused_parameters = cfg.get('find_unused_parameters', False)
         # Sets the `find_unused_parameters` parameter in
         # torch.nn.parallel.DistributedDataParallel
+        if cfg.device == 'npu':
+            current_device = torch.npu.current_device()
+        else:
+            current_device = torch.cuda.current_device()
         model = wrap_distributed_model(
             model,
             cfg.device,
-            device_ids=[torch.cuda.current_device()],
+            device_ids=[current_device],
             broadcast_buffers=False,
             find_unused_parameters=find_unused_parameters)
     else:
@@ -173,6 +177,10 @@ def train_model(model,
 
     # fp16 setting
     fp16_cfg = cfg.get('fp16', None)
+
+    if fp16_cfg is None and device == 'npu':
+        fp16_cfg = {'loss_scale': 'dynamic'}
+
     if fp16_cfg is not None:
         if device == 'ipu':
             from mmcv.device.ipu import IPUFp16OptimizerHook

mmcls/datasets/samplers/distributed_sampler.py

@@ -1,5 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmcv.device.utils import IS_NPU_AVAILABLE
 from torch.utils.data import DistributedSampler as _DistributedSampler
 
 from mmcls.core.utils import sync_random_seed
@@ -30,7 +31,11 @@ def __init__(self,
         # in the same order based on the same seed. Then different ranks
         # could use different indices to select non-overlapped data from the
         # same data list.
-        self.seed = sync_random_seed(seed)
+        if IS_NPU_AVAILABLE:
+            device = 'npu'
+        else:
+            device = 'cuda'
+        self.seed = sync_random_seed(seed, device)
 
     def __iter__(self):
         # deterministically shuffle based on epoch

mmcls/utils/distribution.py

@@ -16,7 +16,10 @@ def wrap_non_distributed_model(model, device='cuda', dim=0, *args, **kwargs):
     Returns:
         model(nn.Module): the model to be parallelized.
     """
-    if device == 'cuda':
+    if device == 'npu':
+        from mmcv.device.npu import NPUDataParallel
+        model = NPUDataParallel(model.npu(), dim=dim, *args, **kwargs)
+    elif device == 'cuda':
         from mmcv.parallel import MMDataParallel
         model = MMDataParallel(model.cuda(), dim=dim, *args, **kwargs)
     elif device == 'cpu':
@@ -49,7 +52,10 @@ def wrap_distributed_model(model, device='cuda', *args, **kwargs):
         .. [1] https://pytorch.org/docs/stable/generated/torch.nn.parallel.
                DistributedDataParallel.html
     """
-    if device == 'cuda':
+    if device == 'npu':
+        from mmcv.device.npu import NPUDistributedDataParallel
+        model = NPUDistributedDataParallel(model.npu(), *args, **kwargs)
+    elif device == 'cuda':
         from mmcv.parallel import MMDistributedDataParallel
         model = MMDistributedDataParallel(model.cuda(), *args, **kwargs)
     else: