Improve default launch device for train

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,32 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 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"