feat: improve launch device

Author: Noé Pion

nerfstudio/configs/base_config.py

@@ -68,8 +68,8 @@ 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 to use for training"""
+    device_type: Literal["cpu", "cuda", "mps"] | None = None
+    """device type to use for training. If none set, script will do its best to set the value."""
 
 
 @dataclass

nerfstudio/scripts/train.py

@@ -62,6 +62,7 @@
 from nerfstudio.configs.method_configs import AnnotatedBaseConfigUnion
 from nerfstudio.engine.trainer import TrainerConfig
 from nerfstudio.utils import comms, profiler
+from nerfstudio.utils.best_device import get_best_device
 from nerfstudio.utils.rich_utils import CONSOLE
 
 DEFAULT_TIMEOUT = timedelta(minutes=30)
@@ -238,6 +239,12 @@ def main(config: TrainerConfig) -> None:
         CONSOLE.log(f"Loading pre-set config from: {config.load_config}")
         config = yaml.load(config.load_config.read_text(), Loader=yaml.Loader)
 
+    if not hasattr(config.machine, "device_type") or config.machine.device_type is None:
+        device_type, reason = get_best_device()
+        config.machine.device_type = device_type
+        CONSOLE.log(f"No device specified: {reason}")
+        CONSOLE.log("You can force a certain device type with --machine.device_type [cuda|mps|cpu]")
+
     config.set_timestamp()
 
     # print and save config

nerfstudio/utils/best_device.py

@@ -0,0 +1,18 @@
+from typing import Literal, Tuple
+
+import torch
+
+
+def get_best_device() -> Tuple[Literal["cpu", "cuda", "mps"], str]:
+    """Determine the best available device to run nerfstudio inference.
+
+    Returns:
+        tuple: (device_type, reason) where device_type is the selected device and
+        reason is an explanation of why it was chosen
+    """
+    if torch.cuda.is_available():
+        return "cuda", "CUDA GPU available - using for optimal performance"
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return "mps", "Apple Metal (MPS) available - using for accelerated performance"
+    else:
+        return "cpu", "No GPU/MPS detected - falling back to CPU"