XPU code update + PyTorch upgrade (#4048)

* draft implementation

* remove unnecessary file

* remove unnecessary directory

* update cuda oriented code

* fix typo

* separate device util files

* add missing file

* fix bug

* use bf16 when mixed precision

* add warn

* remove code to disable xpu

* enable all models in the perf test

* fix linter

* remove warning

* added missing head

* add installation to toml for xpu

* added fixes for few tasks

* fix linter, add device to tests

* remove redundant test

* remove outdated piece of test

* fix linter

* rever MaskDino back

* fix unit test

* installation for XPU is verified

* fix unit tests

* update guide

---------

Co-authored-by: kprokofi <[email protected]>

Author: eunwoosh

docs/source/guide/get_started/installation.rst

@@ -79,71 +79,26 @@ Install OpenVINO™ Training Extensions for users (XPU devices)
 *************************************************************
 
 1. Install OpenVINO™ Training Extensions
-package:
-
-* A local source in development mode
+from source to use XPU functionality.
 
 .. tab-set::
 
-    .. tab-item:: PyPI
-
-        .. code-block:: shell
-
-            pip install otx
-
-    .. tab-item:: Source
+    .. tab-item:: Minimum requirements
 
         .. code-block:: shell
 
             # Clone the training_extensions repository with the following command:
             git clone https://github.com/openvinotoolkit/training_extensions.git
             cd training_extensions
 
-            # Set up a virtual environment.
-            python -m venv .otx
-            source .otx/bin/activate
-
-            pip install -e .
-
-2. Install Intel Extensions for Pytorch & Requirements
-for training according to your system environment.
-
-.. tab-set::
-
-    .. tab-item:: Minimum requirements
-
-        .. code-block:: shell
-
-            pip install '.[xpu]' --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/
-
-        .. note::
-
-            Models from mmlab are not available for this environment. If you want to use mmlab models, you must install them with Full Requirements.
-            Also, some tasks may not be supported by minimum requirements.
+            pip install -e '.[base]' --extra-index-url https://download.pytorch.org/whl/test/xpu
 
-    .. tab-item:: Full Requirements
-
-        .. code-block:: shell
-            python -m pip install torch==2.1.0.post2 torchvision==0.16.0.post2 torchaudio==2.1.0.post2 intel-extension-for-pytorch==2.1.30+xpu oneccl_bind_pt==2.1.300+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/
-            git clone https://github.com/open-mmlab/mmcv
-            cd mmcv
-            git checkout v2.1.0
-            MMCV_WITH_OPS=1 pip install -e .
-            cd ..
-            otx install -v --do-not-install-torch
-
-[Optional] Refer to the `Intel® Extension for PyTorch documentation guide <https://intel.github.io/intel-extension-for-pytorch/index.html#installation?platform=gpu&version=v2.1.10%2bxpu&os=linux%2fwsl2&package=pip>`_
-
-3. Activate OneAPI environment
-and export required IPEX system variables
-
-.. code-block:: shell
+.. note::
 
-    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
+    Please, refer to the `PyTorch official documentation guide <https://pytorch.org/docs/stable/notes/get_start_xpu.html>`_
+    to install prerequisites and resolve possible issues.
 
-3. Once the package is installed in the virtual environment, you can use full
+2. Once the package is installed in the virtual environment, you can use full
 OpenVINO™ Training Extensions command line functionality.
 
 .. code-block:: shell

pyproject.toml

@@ -72,25 +72,8 @@ docs = [
     "nbsphinx",
 ]
 
-xpu = [
-    "torch==2.1.0.post2",
-    "torchvision==0.16.0.post2",
-    "intel-extension-for-pytorch==2.1.30+xpu",
-    "oneccl_bind_pt==2.1.300+xpu",
-    "lightning==2.2",
-    "pytorchcv==0.0.67",
-    "timm==1.0.3",
-    "openvino==2024.4",
-    "openvino-dev==2024.4",
-    "openvino-model-api==0.2.5",
-    "onnx==1.17.0",
-    "onnxconverter-common==1.14.0",
-    "nncf==2.13.0",
-    "anomalib[core]==1.1.0",
-]
-
 base = [
-    "torch==2.4.0",
+    "torch==2.5.1",
     "lightning==2.4.0",
     "pytorchcv==0.0.67",
     "timm==1.0.3",

src/otx/algo/__init__.py

@@ -5,12 +5,10 @@
 
 from . import (
     accelerators,
-    plugins,
     strategies,
 )
 
 __all__ = [
     "strategies",
     "accelerators",
-    "plugins",
 ]

src/otx/algo/accelerators/xpu.py

@@ -10,7 +10,7 @@
 from lightning.pytorch.accelerators import AcceleratorRegistry
 from lightning.pytorch.accelerators.accelerator import Accelerator
 
-from otx.utils.utils import is_xpu_available
+from otx.utils.device import is_xpu_available
 
 
 class XPUAccelerator(Accelerator):

src/otx/algo/common/losses/cross_focal_loss.py

@@ -7,6 +7,7 @@
 
 import torch
 import torch.nn.functional
+from otx.utils.device import get_available_device
 from torch import Tensor, nn
 from torch.amp import custom_fwd
 
@@ -79,7 +80,7 @@ def __init__(
 
         self.cls_criterion = cross_sigmoid_focal_loss
 
-    @custom_fwd(cast_inputs=torch.float32, device_type="cuda")
+    @custom_fwd(device_type=get_available_device(), cast_inputs=torch.float32)
     def forward(
         self,
         pred: Tensor,

src/otx/algo/common/utils/samplers/base_sampler.py

@@ -14,6 +14,7 @@
 import numpy as np
 import torch
 from otx.algo.common.utils.structures import AssignResult, SamplingResult
+from otx.utils.device import is_xpu_available
 
 if TYPE_CHECKING:
     from otx.algo.utils.mmengine_utils import InstanceData
@@ -218,7 +219,13 @@ def random_choice(self, gallery: torch.Tensor | np.ndarray | list, num: int) ->
             raise ValueError(msg)
 
         is_tensor = isinstance(gallery, torch.Tensor)
-        device = torch.cuda.current_device() if torch.cuda.is_available() else "cpu"
+        if is_xpu_available():
+            device = torch.xpu.current_device()
+        elif torch.cuda.is_available():
+            device = torch.cuda.current_device()
+        else:
+            device = "cpu"
+
         _gallery: torch.Tensor = torch.tensor(gallery, dtype=torch.long, device=device) if not is_tensor else gallery
         perm = torch.randperm(_gallery.numel())[:num].to(device=_gallery.device)
         rand_inds = _gallery[perm]

src/otx/algo/detection/atss.py

@@ -27,7 +27,6 @@
 if TYPE_CHECKING:
     from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
     from torch import nn
-    from typing_extensions import Self
 
     from otx.core.metrics import MetricCallable
     from otx.core.schedulers import LRSchedulerListCallable
@@ -195,11 +194,3 @@ def _exporter(self) -> OTXModelExporter:
     def load_from_otx_v1_ckpt(self, state_dict: dict, add_prefix: str = "model.") -> dict:
         """Load the previous OTX ckpt according to OTX2.0."""
         return OTXv1Helper.load_det_ckpt(state_dict, add_prefix)
-
-    def to(self, *args, **kwargs) -> Self:
-        """Return a model with specified device."""
-        ret = super().to(*args, **kwargs)
-        if self.model_name == "atss_resnext101" and self.device.type == "xpu":
-            msg = f"{type(self).__name__} doesn't support XPU."
-            raise RuntimeError(msg)
-        return ret

src/otx/algo/detection/utils/assigners/sim_ota_assigner.py

@@ -90,6 +90,8 @@ def assign(
         priors = pred_instances.priors  # type: ignore[attr-defined]
         num_bboxes = decoded_bboxes.size(0)
 
+        device = gt_bboxes.device.type
+
         # assign 0 by default
         assigned_gt_inds = decoded_bboxes.new_full((num_bboxes,), 0, dtype=torch.long)
         if num_gt == 0 or num_bboxes == 0:
@@ -117,7 +119,7 @@ def assign(
 
         valid_pred_scores = valid_pred_scores.unsqueeze(1).repeat(1, num_gt, 1)
         # disable AMP autocast and calculate BCE with FP32 to avoid overflow
-        with torch.cuda.amp.autocast(enabled=False):
+        with torch.amp.autocast(device_type=device, enabled=False):
             cls_cost = (
                 F.binary_cross_entropy(
                     valid_pred_scores.to(dtype=torch.float32),

src/otx/algo/keypoint_detection/heads/rtmcc_head.py

@@ -192,7 +192,7 @@ def loss(self, x: tuple[Tensor], entity: KeypointDetBatchDataEntity) -> dict:
             mask=mask,
         )
 
-        loss_pose = -1 * torch.tensor(avg_acc, device=device)
+        loss_pose = -1 * torch.tensor(avg_acc, device=device, dtype=torch.float32)
         losses.update(loss_pose=loss_pose)
 
         return losses