added Intel GPU support for training and inference

nerfstudio/configs/base_config.py

@@ -69,7 +69,7 @@ class MachineConfig(PrintableConfig):
     """current machine's rank (for DDP)"""
     dist_url: str = "auto"
     """distributed connection point (for DDP)"""
-    device_type: Literal["cpu", "cuda", "mps"] = "cuda"
+    device_type: Literal["cpu", "cuda", "mps", "xpu"] = "cuda"
     """device type to use for training"""
 
 

nerfstudio/engine/trainer.py

@@ -42,6 +42,7 @@
 from rich.panel import Panel
 from rich.table import Table
 from torch.cuda.amp.grad_scaler import GradScaler
+import importlib
 
 TRAIN_INTERATION_OUTPUT = Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]
 TORCH_DEVICE = str
@@ -115,12 +116,14 @@ def __init__(self, config: TrainerConfig, local_rank: int = 0, world_size: int =
         self.device: TORCH_DEVICE = config.machine.device_type
         if self.device == "cuda":
             self.device += f":{local_rank}"
+        elif self.device == "xpu":
+            importlib.import_module("intel_extension_for_pytorch")
         self.mixed_precision: bool = self.config.mixed_precision
         self.use_grad_scaler: bool = self.mixed_precision or self.config.use_grad_scaler
         self.training_state: Literal["training", "paused", "completed"] = "training"
         self.gradient_accumulation_steps: int = self.config.gradient_accumulation_steps
 
-        if self.device == "cpu":
+        if self.device == "cpu" or self.device == "xpu":
             self.mixed_precision = False
             CONSOLE.print("Mixed precision is disabled for CPU training.")
         self._start_step: int = 0
@@ -460,13 +463,13 @@ def train_iteration(self, step: int) -> TRAIN_INTERATION_OUTPUT:
         """
 
         self.optimizers.zero_grad_all()
-        cpu_or_cuda_str: str = self.device.split(":")[0]
-        cpu_or_cuda_str = "cpu" if cpu_or_cuda_str == "mps" else cpu_or_cuda_str
+        cpu_or_gpu_str: str = self.device.split(":")[0]
+        cpu_or_gpu_str = "cpu" if cpu_or_gpu_str == "mps" else cpu_or_gpu_str
         assert (
             self.gradient_accumulation_steps > 0
         ), f"gradient_accumulation_steps must be > 0, not {self.gradient_accumulation_steps}"
         for _ in range(self.gradient_accumulation_steps):
-            with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision):
+            with torch.autocast(device_type=cpu_or_gpu_str, enabled=self.mixed_precision):
                 _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step)
                 loss = functools.reduce(torch.add, loss_dict.values())
                 loss /= self.gradient_accumulation_steps

nerfstudio/scripts/train.py

@@ -109,7 +109,7 @@ def _distributed_worker(
     dist_url: str,
     config: TrainerConfig,
     timeout: timedelta = DEFAULT_TIMEOUT,
-    device_type: Literal["cpu", "cuda", "mps"] = "cuda",
+    device_type: Literal["cpu", "cuda", "mps", "xpu"] = "cuda",
 ) -> Any:
     """Spawned distributed worker that handles the initialization of process group and handles the
        training process on multiple processes.
@@ -165,7 +165,7 @@ def launch(
     dist_url: str = "auto",
     config: Optional[TrainerConfig] = None,
     timeout: timedelta = DEFAULT_TIMEOUT,
-    device_type: Literal["cpu", "cuda", "mps"] = "cuda",
+    device_type: Literal["cpu", "cuda", "mps", "xpu"] = "cuda",
 ) -> None:
     """Function that spawns multiple processes to call on main_func
 

nerfstudio/utils/eval_utils.py

@@ -16,6 +16,7 @@
 Evaluation utils
 """
 from __future__ import annotations
+import importlib
 
 import os
 import sys
@@ -101,7 +102,15 @@ def eval_setup(
     config.load_dir = config.get_checkpoint_dir()
 
     # setup pipeline (which includes the DataManager)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if torch.cuda.is_available():
+        device_type = "cuda"
+    else:
+        try:
+            importlib.import_module("intel_extension_for_pytorch")
+            device_type = "xpu" if torch.xpu.is_available() else "cpu"
+        except Exception:
+            device_type = "cpu"
+    device = torch.device(device_type)
     pipeline = config.pipeline.setup(device=device, test_mode=test_mode)
     assert isinstance(pipeline, Pipeline)
     pipeline.eval()