Add support of the Intel dGPUs devices

src/otx/algorithms/anomaly/adapters/__init__.py

@@ -13,3 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions
 # and limitations under the License.
+
+
+from .anomalib.accelerators.xpu import XPUAccelerator  # noqa: F401
+from .anomalib.strategies import SingleXPUStrategy  # noqa: F401

src/otx/algorithms/anomaly/adapters/anomalib/accelerators/__init__.py

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

src/otx/algorithms/anomaly/adapters/anomalib/accelerators/xpu.py

@@ -0,0 +1,60 @@
+"""Lightning accelerator for XPU device."""
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+
+from typing import Any, Dict, Union
+
+import torch
+from pytorch_lightning.accelerators import AcceleratorRegistry
+from pytorch_lightning.accelerators.accelerator import Accelerator
+
+from otx.algorithms.common.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":
+            raise RuntimeError(f"Device should be xpu, got {device} instead")
+
+        torch.xpu.set_device(device)
+
+    @staticmethod
+    def parse_devices(devices: Any) -> Any:
+        """Parses devices for multi-GPU training."""
+        if isinstance(devices, list):
+            return devices
+        return [devices]
+
+    @staticmethod
+    def get_parallel_devices(devices: Any) -> Any:
+        """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: Union[str, torch.device]) -> Dict[str, Any]:
+        """Returns XPU devices stats."""
+        return {}
+
+    def teardown(self) -> None:
+        """Cleans-up XPU-related resources."""
+        pass
+
+
+AcceleratorRegistry.register(
+    XPUAccelerator.accelerator_name, XPUAccelerator, description="Accelerator supports XPU devices"
+)

src/otx/algorithms/anomaly/adapters/anomalib/plugins/__init__.py

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

src/otx/algorithms/anomaly/adapters/anomalib/plugins/xpu_precision.py

@@ -0,0 +1,109 @@
+"""Plugin for mixed-precision training on XPU."""
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, Generator, Optional, Union
+
+import pytorch_lightning as pl
+import torch
+from lightning_fabric.utilities.types import Optimizable
+from pytorch_lightning.plugins.precision.precision_plugin import PrecisionPlugin
+from pytorch_lightning.utilities import GradClipAlgorithmType
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from torch import Tensor
+from torch.optim import LBFGS, Optimizer
+
+
+class MixedPrecisionXPUPlugin(PrecisionPlugin):
+    """Plugin for Automatic Mixed Precision (AMP) training with ``torch.xpu.autocast``.
+
+    Args:
+        scaler: An optional :class:`torch.cuda.amp.GradScaler` to use.
+    """
+
+    def __init__(self, scaler: Optional[Any] = None) -> None:
+        self.scaler = scaler
+
+    def pre_backward(self, tensor: Tensor, module: "pl.LightningModule") -> Tensor:
+        """Apply grad scaler before backward."""
+        if self.scaler is not None:
+            tensor = self.scaler.scale(tensor)
+        return super().pre_backward(tensor, module)
+
+    def optimizer_step(  # type: ignore[override]
+        self,
+        optimizer: Optimizable,
+        model: "pl.LightningModule",
+        optimizer_idx: int,
+        closure: Callable[[], Any],
+        **kwargs: Any,
+    ) -> Any:
+        """Make an optimizer step using scaler if it was passed."""
+        if self.scaler is None:
+            # skip scaler logic, as bfloat16 does not require scaler
+            return super().optimizer_step(
+                optimizer, model=model, optimizer_idx=optimizer_idx, closure=closure, **kwargs
+            )
+        if isinstance(optimizer, LBFGS):
+            raise MisconfigurationException(
+                f"Native AMP and the LBFGS optimizer are not compatible (optimizer {optimizer_idx})."
+            )
+        closure_result = closure()
+
+        if not _optimizer_handles_unscaling(optimizer):
+            # Unscaling needs to be performed here in case we are going to apply gradient clipping.
+            # Optimizers that perform unscaling in their `.step()` method are not supported (e.g., fused Adam).
+            # Note: `unscale` happens after the closure is executed, but before the `on_before_optimizer_step` hook.
+            self.scaler.unscale_(optimizer)
+
+        self._after_closure(model, optimizer, optimizer_idx)
+        skipped_backward = closure_result is None
+        # in manual optimization, the closure does not return a value
+        if not model.automatic_optimization or not skipped_backward:
+            # note: the scaler will skip the `optimizer.step` if nonfinite gradients are found
+            step_output = self.scaler.step(optimizer, **kwargs)
+            self.scaler.update()
+            return step_output
+        return closure_result
+
+    def clip_gradients(
+        self,
+        optimizer: Optimizer,
+        clip_val: Union[int, float] = 0.0,
+        gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
+    ) -> None:
+        """Handle grad clipping with scaler."""
+        if clip_val > 0 and _optimizer_handles_unscaling(optimizer):
+            raise RuntimeError(
+                f"The current optimizer, {type(optimizer).__qualname__}, does not allow for gradient clipping"
+                " because it performs unscaling of gradients internally. HINT: Are you using a 'fused' optimizer?"
+            )
+        super().clip_gradients(optimizer=optimizer, clip_val=clip_val, gradient_clip_algorithm=gradient_clip_algorithm)
+
+    @contextmanager
+    def forward_context(self) -> Generator[None, None, None]:
+        """Enable autocast context."""
+        with torch.xpu.autocast(True):
+            yield
+
+    def state_dict(self) -> Dict[str, Any]:
+        """Returns state dict of the plugin."""
+        if self.scaler is not None:
+            return self.scaler.state_dict()
+        return {}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        """Loads state dict to the plugin."""
+        if self.scaler is not None:
+            self.scaler.load_state_dict(state_dict)
+
+
+def _optimizer_handles_unscaling(optimizer: Any) -> bool:
+    """Determines if a PyTorch optimizer handles unscaling gradients in the step method ratherthan through the scaler.
+
+    Since, the current implementation of this function checks a PyTorch internal variable on the optimizer, the return
+    value will only be reliable for built-in PyTorch optimizers.
+    """
+    return getattr(optimizer, "_step_supports_amp_scaling", False)

src/otx/algorithms/anomaly/adapters/anomalib/strategies/__init__.py

@@ -0,0 +1,8 @@
+"""Lightning strategy for single XPU device."""
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+
+from .xpu_single import SingleXPUStrategy
+
+__all__ = ["SingleXPUStrategy"]

src/otx/algorithms/anomaly/adapters/anomalib/strategies/xpu_single.py

@@ -0,0 +1,60 @@
+"""Lightning strategy for single XPU device."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+
+from typing import Optional
+
+import pytorch_lightning as pl
+import torch
+from lightning_fabric.plugins import CheckpointIO
+from lightning_fabric.utilities.types import _DEVICE
+from pytorch_lightning.plugins.precision import PrecisionPlugin
+from pytorch_lightning.strategies import StrategyRegistry
+from pytorch_lightning.strategies.single_device import SingleDeviceStrategy
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+
+from otx.algorithms.common.utils.utils import is_xpu_available
+
+
+class SingleXPUStrategy(SingleDeviceStrategy):
+    """Strategy for training on single XPU device."""
+
+    strategy_name = "xpu_single"
+
+    def __init__(
+        self,
+        device: _DEVICE = "xpu:0",
+        accelerator: Optional["pl.accelerators.Accelerator"] = None,
+        checkpoint_io: Optional[CheckpointIO] = None,
+        precision_plugin: Optional[PrecisionPlugin] = None,
+    ):
+        if not is_xpu_available():
+            raise MisconfigurationException("`SingleXPUStrategy` requires XPU devices to run")
+
+        super().__init__(
+            accelerator=accelerator,
+            device=device,
+            checkpoint_io=checkpoint_io,
+            precision_plugin=precision_plugin,
+        )
+
+    @property
+    def is_distributed(self) -> bool:
+        """Returns true if the strategy supports distributed training."""
+        return False
+
+    def setup_optimizers(self, trainer: "pl.Trainer") -> None:
+        """Sets up optimizers."""
+        super().setup_optimizers(trainer)
+        if len(self.optimizers) != 1:  # type: ignore
+            raise RuntimeError("XPU strategy doesn't support multiple optimizers")
+        model, optimizer = torch.xpu.optimize(trainer.model, optimizer=self.optimizers[0])  # type: ignore
+        self.optimizers = [optimizer]
+        trainer.model = model
+
+
+StrategyRegistry.register(
+    SingleXPUStrategy.strategy_name, SingleXPUStrategy, description="Strategy that enables training on single XPU"
+)

src/otx/algorithms/anomaly/tasks/train.py

@@ -30,6 +30,8 @@
 
 from otx.algorithms.anomaly.adapters.anomalib.callbacks import IterationTimer, ProgressCallback
 from otx.algorithms.anomaly.adapters.anomalib.data import OTXAnomalyDataModule
+from otx.algorithms.anomaly.adapters.anomalib.plugins.xpu_precision import MixedPrecisionXPUPlugin
+from otx.algorithms.common.utils.utils import is_xpu_available
 from otx.api.entities.datasets import DatasetEntity
 from otx.api.entities.model import ModelEntity
 from otx.api.entities.train_parameters import TrainParameters
@@ -90,7 +92,20 @@
             IterationTimer(on_step=False),
         ]
 
-        self.trainer = Trainer(**config.trainer, logger=CSVLogger(self.project_path, name=""), callbacks=callbacks)
+        plugins = []
+        if config.trainer.plugins is not None:
+            plugins.extend(config.trainer.plugins)
+        config.trainer.pop("plugins")
+
+        if is_xpu_available():
+            config.trainer.strategy = "xpu_single"
+            config.trainer.accelerator = "xpu"
+            if config.trainer.precision == 16:
+                plugins.append(MixedPrecisionXPUPlugin())
+
+        self.trainer = Trainer(
+            **config.trainer, logger=CSVLogger(self.project_path, name=""), callbacks=callbacks, plugins=plugins
+        )
         self.trainer.fit(model=self.model, datamodule=datamodule)
 
         self.save_model(output_model)

src/otx/algorithms/classification/adapters/mmcls/apis/__init__.py

@@ -0,0 +1,8 @@
+"""Adapters of classification - mmcls."""
+
+# Copyright (C) 2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .train import train_model
+
+__all__ = ["train_model"]