Enable training on XPU devices in OTX2.0

docs/source/guide/get_started/installation.rst

@@ -11,9 +11,9 @@ The current version of OpenVINO™ Training Extensions was tested in the followi
 - Python >= 3.10
 
 
-***********************************************
-Install OpenVINO™ Training Extensions for users
-***********************************************
+**********************************************************
+Install OpenVINO™ Training Extensions for users (CUDA/CPU)
+**********************************************************
 
 1. Install OpenVINO™ Training Extensions package:
 
@@ -57,6 +57,68 @@ Install OpenVINO™ Training Extensions for users
 3. Once the package is installed in the virtual environment, you can use full
 OpenVINO™ Training Extensions command line functionality.
 
+*************************************************************
+Install OpenVINO™ Training Extensions for users (XPU devices)
+*************************************************************
+
+1. Follow the first two steps from above instructions
+on cloning the repository and creating a virtual environment.
+
+2. Install Intel Extensions For Pytorch (IPEX).
+Follow the `official documentation <https://intel.github.io/intel-extension-for-pytorch/index.html#installation?platform=gpu&version=v2.1.10%2Bxpu>`_ to install prerequisites such as OneAPI and proper drivers.
+
+.. code-block:: shell
+
+    python -m pip install torch==2.1.0a0 torchvision==0.16.0a0 torchaudio==2.1.0a0 intel-extension-for-pytorch==2.1.10+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/
+
+3. Install MMCV.
+It is required to install mmcv from source to properly build it with IPEX.
+
+.. code-block:: shell
+
+    git clone https://github.com/open-mmlab/mmcv
+    cd mmcv
+    git checkout v2.1.0
+    MMCV_WITH_OPS=1 pip install -e .
+
+4. Install OpenVINO™ Training Extensions
+package from either:
+
+* A local source in development mode
+
+.. code-block:: shell
+
+    pip install -e .
+
+* PyPI
+
+.. code-block:: shell
+
+    pip install otx
+
+5. Install requirements for training
+excluding Pytorch.
+
+.. code-block:: shell
+
+    otx install -v --do-not-install-torch
+
+6. Activate OneAPI environment
+and export required IPEX system variables
+
+.. code-block:: shell
+
+    source /path/to/intel/oneapi/setvars.sh
+    export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libstdc++.so.6.0.30
+    export IPEX_FP32_MATH_MODE=TF32
+
+7. Once the package is installed in the virtual environment, you can use full
+OpenVINO™ Training Extensions command line functionality.
+
+.. code-block:: shell
+
+    otx --help
+
 ****************************************************
 Install OpenVINO™ Training Extensions for developers
 ****************************************************

src/otx/algo/__init__.py

@@ -3,6 +3,24 @@
 #
 """Module for OTX custom algorithms, e.g., model, losses, hook, etc..."""
 
-from . import action_classification, classification, detection, segmentation, visual_prompting
+from . import (
+    accelerators,
+    action_classification,
+    classification,
+    detection,
+    plugins,
+    segmentation,
+    strategies,
+    visual_prompting,
+)
 
-__all__ = ["action_classification", "classification", "detection", "segmentation", "visual_prompting"]
+__all__ = [
+    "action_classification",
+    "classification",
+    "detection",
+    "segmentation",
+    "visual_prompting",
+    "strategies",
+    "accelerators",
+    "plugins",
+]

src/otx/algo/accelerators/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+"""Lightning accelerator for XPU device."""
+
+from .xpu import XPUAccelerator
+
+__all__ = ["XPUAccelerator"]

src/otx/algo/accelerators/xpu.py

@@ -0,0 +1,88 @@
+"""Lightning accelerator for XPU device."""
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+from __future__ import annotations
+
+from typing import Any, Union
+
+import numpy as np
+import torch
+from lightning.pytorch.accelerators import AcceleratorRegistry
+from lightning.pytorch.accelerators.accelerator import Accelerator
+from mmcv.ops.nms import NMSop
+from mmcv.ops.roi_align import RoIAlign
+from mmengine.structures import instance_data
+
+from otx.algo.detection.utils import monkey_patched_nms, monkey_patched_roi_align
+from otx.utils.utils import is_xpu_available
+
+
+class XPUAccelerator(Accelerator):
+    """Support for a XPU, optimized for large-scale machine learning."""
+
+    accelerator_name = "xpu"
+
+    def setup_device(self, device: torch.device) -> None:
+        """Sets up the specified device."""
+        if device.type != "xpu":
+            msg = f"Device should be xpu, got {device} instead"
+            raise RuntimeError(msg)
+
+        torch.xpu.set_device(device)
+        self.patch_packages_xpu()
+
+    @staticmethod
+    def parse_devices(devices: str | list | torch.device) -> list:
+        """Parses devices for multi-GPU training."""
+        if isinstance(devices, list):
+            return devices
+        return [devices]
+
+    @staticmethod
+    def get_parallel_devices(devices: list) -> list[torch.device]:
+        """Generates a list of parrallel devices."""
+        return [torch.device("xpu", idx) for idx in devices]
+
+    @staticmethod
+    def auto_device_count() -> int:
+        """Returns number of XPU devices available."""
+        return torch.xpu.device_count()
+
+    @staticmethod
+    def is_available() -> bool:
+        """Checks if XPU available."""
+        return is_xpu_available()
+
+    def get_device_stats(self, device: str | torch.device) -> dict[str, Any]:
+        """Returns XPU devices stats."""
+        return {}
+
+    def teardown(self) -> None:
+        """Cleans-up XPU-related resources."""
+        self.revert_packages_xpu()
+
+    def patch_packages_xpu(self) -> None:
+        """Patch packages when xpu is available."""
+        # patch instance_data from mmengie
+        long_type_tensor = Union[torch.LongTensor, torch.xpu.LongTensor]
+        bool_type_tensor = Union[torch.BoolTensor, torch.xpu.BoolTensor]
+        instance_data.IndexType = Union[str, slice, int, list, long_type_tensor, bool_type_tensor, np.ndarray]
+
+        # patch nms and roi_align
+        self._nms_op_forward = NMSop.forward
+        self._roi_align_forward = RoIAlign.forward
+        NMSop.forward = monkey_patched_nms
+        RoIAlign.forward = monkey_patched_roi_align
+
+    def revert_packages_xpu(self) -> None:
+        """Revert packages when xpu is available."""
+        NMSop.forward = self._nms_op_forward
+        RoIAlign.forward = self._roi_align_forward
+
+
+AcceleratorRegistry.register(
+    XPUAccelerator.accelerator_name,
+    XPUAccelerator,
+    description="Accelerator supports XPU devices",
+)

src/otx/algo/detection/utils/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+"""utils for detection task."""
+
+from .mmcv_patched_ops import monkey_patched_nms, monkey_patched_roi_align
+
+__all__ = ["monkey_patched_nms", "monkey_patched_roi_align"]

src/otx/algo/detection/utils/mmcv_patched_ops.py

@@ -0,0 +1,73 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+"""utils for detection task."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from mmcv.utils import ext_loader
+from torchvision.ops import nms as tv_nms
+from torchvision.ops import roi_align as tv_roi_align
+
+if TYPE_CHECKING:
+    from mmcv.ops.nms import NMSop
+    from mmcv.ops.roi_align import RoIAlign
+
+ext_module = ext_loader.load_ext("_ext", ["nms", "softnms", "nms_match", "nms_rotated", "nms_quadri"])
+
+
+def monkey_patched_nms(
+    ctx: NMSop,
+    bboxes: torch.Tensor,
+    scores: torch.Tensor,
+    iou_threshold: float,
+    offset: float,
+    score_threshold: float,
+    max_num: int,
+) -> torch.Tensor:
+    """Runs MMCVs NMS with torchvision.nms, or forces NMS from MMCV to run on CPU."""
+    _ = ctx
+    is_filtering_by_score = score_threshold > 0
+    if is_filtering_by_score:
+        valid_mask = scores > score_threshold
+        bboxes, scores = bboxes[valid_mask], scores[valid_mask]
+        valid_inds = torch.nonzero(valid_mask, as_tuple=False).squeeze(dim=1)
+
+    if bboxes.dtype == torch.bfloat16:
+        bboxes = bboxes.to(torch.float32)
+    if scores.dtype == torch.bfloat16:
+        scores = scores.to(torch.float32)
+
+    if offset == 0:
+        inds = tv_nms(bboxes, scores, float(iou_threshold))
+    else:
+        device = bboxes.device
+        bboxes = bboxes.to("cpu")
+        scores = scores.to("cpu")
+        inds = ext_module.nms(bboxes, scores, iou_threshold=float(iou_threshold), offset=offset)
+        bboxes = bboxes.to(device)
+        scores = scores.to(device)
+
+    if max_num > 0:
+        inds = inds[:max_num]
+    if is_filtering_by_score:
+        inds = valid_inds[inds]
+    return inds
+
+
+def monkey_patched_roi_align(self: RoIAlign, _input: torch.Tensor, rois: torch.Tensor) -> torch.Tensor:
+    """Replaces MMCVs roi align with the one from torchvision.
+
+    Args:
+        self: patched instance
+        _input: NCHW images
+        rois: Bx5 boxes. First column is the index into N. The other 4 columns are xyxy.
+    """
+    if "aligned" in tv_roi_align.__code__.co_varnames:
+        return tv_roi_align(_input, rois, self.output_size, self.spatial_scale, self.sampling_ratio, self.aligned)
+    if self.aligned:
+        rois -= rois.new_tensor([0.0] + [0.5 / self.spatial_scale] * 4)
+    return tv_roi_align(_input, rois, self.output_size, self.spatial_scale, self.sampling_ratio)

src/otx/algo/plugins/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+"""Plugin for mixed-precision training on XPU."""
+
+from .xpu_precision import MixedPrecisionXPUPlugin
+
+__all__ = ["MixedPrecisionXPUPlugin"]