diff --git a/mmengine/device/utils.py b/mmengine/device/utils.py index 1a31af549b..44e92f715b 100644 --- a/mmengine/device/utils.py +++ b/mmengine/device/utils.py @@ -36,6 +36,10 @@ def is_npu_available() -> bool: """Returns True if Ascend PyTorch and npu devices exist.""" try: import torch_npu # noqa: F401 + + # Enable operator support for dynamic shape and + # binary operator support on the NPU. + torch.npu.set_compile_mode(jit_compile=False) except Exception: return False return hasattr(torch, 'npu') and torch.npu.is_available() diff --git a/mmengine/model/base_model/base_model.py b/mmengine/model/base_model/base_model.py index f9316506d8..06bf5b6594 100644 --- a/mmengine/model/base_model/base_model.py +++ b/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('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)) diff --git a/mmengine/model/base_model/data_preprocessor.py b/mmengine/model/base_model/data_preprocessor.py index 17d70e067e..78a55bced1 100644 --- a/mmengine/model/base_model/data_preprocessor.py +++ b/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('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` diff --git a/mmengine/structures/base_data_element.py b/mmengine/structures/base_data_element.py index 042a9df673..7be1ef9044 100644 --- a/mmengine/structures/base_data_element.py +++ b/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."""