add npu device support

mmengine/device/utils.py

@@ -36,6 +36,7 @@ def is_npu_available() -> bool:
     """Returns True if Ascend PyTorch and npu devices exist."""
     try:
         import torch_npu  # noqa: F401
+        torch.npu.set_compile_mode(jit_compile=False)
     except Exception:
         return False
     return hasattr(torch, 'npu') and torch.npu.is_available()

mmengine/model/base_model/base_model.py

@@ -184,6 +184,18 @@ def to(self, *args, **kwargs) -> nn.Module:
         Returns:
             nn.Module: The model itself.
         """
+
+        # Since Torch has not officially merged
+        # the npu-related fields, using the _parse_to function
+        # directly will cause the NPU to not be found.
+        # Here, the input parameters are processed to avoid errors.
+        if args and isinstance(args[0], str) and 'npu' in args[0]:
+            args = tuple([list(args)])[0].replace(  # type: ignore
+                'npu', torch.npu.native_device)
+        if kwargs and 'npu' in str(kwargs.get('device', '')):
+            kwargs['device'] = kwargs['device'].replace(
+                'npu', torch.npu.native_device)
+
         device = torch._C._nn._parse_to(*args, **kwargs)[0]
         if device is not None:
             self._set_device(torch.device(device))

mmengine/model/base_model/data_preprocessor.py

@@ -90,6 +90,18 @@ def to(self, *args, **kwargs) -> nn.Module:
         Returns:
             nn.Module: The model itself.
         """
+
+        # Since Torch has not officially merged
+        # the npu-related fields, using the _parse_to function
+        # directly will cause the NPU to not be found.
+        # Here, the input parameters are processed to avoid errors.
+        if args and isinstance(args[0], str) and 'npu' in args[0]:
+            args = tuple([list(args)])[0].replace(  # type: ignore
+                'npu', torch.npu.native_device)
+        if kwargs and 'npu' in str(kwargs.get('device', '')):
+            kwargs['device'] = kwargs['device'].replace(
+                'npu', torch.npu.native_device)
+
         device = torch._C._nn._parse_to(*args, **kwargs)[0]
         if device is not None:
             self._device = torch.device(device)
@@ -104,6 +116,15 @@ def cuda(self, *args, **kwargs) -> nn.Module:
         self._device = torch.device(torch.cuda.current_device())
         return super().cuda()
 
+    def npu(self, *args, **kwargs) -> nn.Module:
+        """Overrides this method to set the :attr:`device`
+
+        Returns:
+            nn.Module: The model itself.
+        """
+        self._device = torch.device(torch.npu.current_device())
+        return super().npu()
+
     def cpu(self, *args, **kwargs) -> nn.Module:
         """Overrides this method to set the :attr:`device`
 

mmengine/structures/base_data_element.py

@@ -507,6 +507,17 @@ def cuda(self) -> 'BaseDataElement':
                 new_data.set_data(data)
         return new_data
 
+    # Tensor-like methods
+    def npu(self) -> 'BaseDataElement':
+        """Convert all tensors to NPU in data."""
+        new_data = self.new()
+        for k, v in self.items():
+            if isinstance(v, (torch.Tensor, BaseDataElement)):
+                v = v.npu()
+                data = {k: v}
+                new_data.set_data(data)
+        return new_data
+
     # Tensor-like methods
     def detach(self) -> 'BaseDataElement':
         """Detach all tensors in data."""