diff --git a/mmengine/device/utils.py b/mmengine/device/utils.py index 1a31af549be57995eea660fee0f8eb21db9bdc44..44e92f715bf701627990a131520b42a2733e1376 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 f9316506d8613d1f500f9325a6a0aebea3e9e360..06bf5b659459e44160ea7d128a5ad5b138066bf8 100644 --- a/mmengine/model/base_model/base_model.py +++ b/mmengine/model/base_model/base_model.py @@ -184,6 +184,18 @@ class BaseModel(BaseModule): 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 14c3db96c1bcaca73b87dccda4357ca48b57db94..8f02a6c5a4c597cb1d41e0f6550f85a95139c757 100644 --- a/mmengine/model/base_model/data_preprocessor.py +++ b/mmengine/model/base_model/data_preprocessor.py @@ -87,6 +87,18 @@ class BaseDataPreprocessor(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) @@ -101,6 +113,15 @@ class BaseDataPreprocessor(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 042a9df6737306ceb39ed97c73a0818f9a690b83..7be1ef90447a21b11f3a154d4507c8815ef6a537 100644 --- a/mmengine/structures/base_data_element.py +++ b/mmengine/structures/base_data_element.py @@ -507,6 +507,17 @@ class 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."""