updated device types

Author: mhubii

cli/rr_cam_swarm.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+from typing import Union
 
 import cv2
 import numpy as np
@@ -179,7 +180,7 @@ def instantiate_particles(
     eye_min_dist: float,
     eye_max_dist: float,
     angle_interval: float,
-    device: torch.device = torch.device("cuda"),
+    device: Union[torch.device, str] = "cuda",
 ) -> torch.Tensor:
     r"""Instantiate the particles for the optimization randomly under field of view constraints.
     Particles (camera poses) are represented using eye space coordinates (eye, center, angle).
@@ -193,7 +194,7 @@ def instantiate_particles(
         eye_min_dist (float): The minimum distance of the eye from the origin.
         eye_max_dist (float): The maximum distance of the eye from the origin.
         angle_interval (float): The angle interval in which to sample the rotation angle.
-        device (torch.device): The device to instantiate the particles on.
+        device (Union[torch.device, str]): The device to instantiate the particles on.
 
     Returns:
         torch.Tensor: The particles of shape (n_particles, 7).

roboreg/differentiable/kinematics.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Union
 
 import pytorch_kinematics as pk
 import torch
@@ -17,7 +17,7 @@ def __init__(
         urdf: str,
         root_link_name: str,
         end_link_name: str,
-        device: torch.device = torch.device("cuda"),
+        device: Union[torch.device, str] = "cuda",
     ) -> None:
         self._root_link_name = root_link_name
         self._end_link_name = end_link_name
@@ -26,7 +26,7 @@ def __init__(
             root_link_name=self._root_link_name,
             end_link_name=self._end_link_name,
         )
-        self._device = device
+        self._device = torch.device(device) if isinstance(device, str) else device
         self.to(device=self._device)
 
     def _build_serial_chain_from_urdf(
@@ -36,9 +36,9 @@ def _build_serial_chain_from_urdf(
             urdf, end_link_name=end_link_name, root_link_name=root_link_name
         )
 
-    def to(self, device: torch.device) -> None:
+    def to(self, device: Union[torch.device, str]) -> None:
         self._chain.to(device=device)
-        self._device = device
+        self._device = torch.device(device) if isinstance(device, str) else device
 
     def forward_kinematics(self, q: torch.Tensor) -> Dict[str, torch.Tensor]:
         ht_lookup = {

roboreg/differentiable/rendering.py

@@ -1,4 +1,4 @@
-from typing import List, Tuple
+from typing import List, Tuple, Union
 
 import nvdiffrast.torch as dr
 import torch
@@ -13,11 +13,11 @@ class NVDiffRastRenderer:
 
     __slots__ = ["_device", "_ctx"]
 
-    def __init__(self, device: torch.device = torch.device("cuda")) -> None:
+    def __init__(self, device: Union[torch.device, str] = "cuda") -> None:
         super().__init__()
         if not torch.cuda.is_available():
             raise ValueError("CUDA is not available.")
-        self._device = device
+        self._device = torch.device(device) if isinstance(device, str) else device
         self._ctx = dr.RasterizeCudaContext(device=self._device)
 
     def scale_clip_vertices(

roboreg/differentiable/robot.py

@@ -1,3 +1,5 @@
+from typing import Union
+
 import torch
 
 from roboreg.io import URDFParser
@@ -13,12 +15,12 @@ def __init__(
         self,
         mesh_container: TorchMeshContainer,
         kinematics: TorchKinematics,
-        device: torch.device = torch.device("cuda"),
+        device: Union[torch.device, str] = "cuda",
     ) -> None:
         self._mesh_container = mesh_container
         self._kinematics = kinematics
         self._configured_vertices = self.mesh_container.vertices.clone()
-        self._device = device
+        self._device = torch.device(device) if isinstance(device, str) else device
         self.to(device=self._device)
 
     @classmethod
@@ -29,7 +31,7 @@ def from_urdf_parser(
         end_link_name: str,
         collision: bool = False,
         batch_size: int = 1,
-        device: torch.device = torch.device("cuda"),
+        device: Union[torch.device, str] = "cuda",
         target_reduction: float = 0.0,
     ) -> "Robot":
         from roboreg.io import apply_mesh_origins, load_meshes, simplify_meshes
@@ -104,11 +106,11 @@ def configure(
                 ht_root.transpose(-1, -2),
             )
 
-    def to(self, device: torch.device) -> None:
+    def to(self, device: Union[torch.device, str]) -> None:
         self._mesh_container.to(device=device)
         self._kinematics.to(device=device)
         self._configured_vertices = self._configured_vertices.to(device=device)
-        self._device = device
+        self._device = torch.device(device) if isinstance(device, str) else device
 
     @property
     def kinematics(self) -> TorchKinematics:

roboreg/differentiable/structs.py

@@ -34,7 +34,7 @@ def __init__(
         self,
         meshes: Dict[str, Mesh],
         batch_size: int = 1,
-        device: torch.device = torch.device("cuda"),
+        device: Union[torch.device, str] = "cuda",
     ) -> None:
         self._names = []
         self._vertices = []
@@ -47,7 +47,7 @@ def __init__(
         self._upper_face_index_lookup = {}
 
         # populate this container
-        self._populate_container(meshes, device)
+        self._populate_container(meshes, device=device)
 
         # add batch dim
         self._batch_size = batch_size
@@ -56,13 +56,13 @@ def __init__(
         # populate index lookups
         self._populate_index_lookups()
 
-        self._device = device
+        self._device = torch.device(device) if isinstance(device, str) else device
 
     @abc.abstractmethod
     def _populate_container(
         self,
         meshes: Dict[str, Mesh],
-        device: torch.device("cuda"),
+        device: Union[torch.device, str] = "cuda",
     ) -> None:
         offset = 0
         for name, mesh in meshes.items():
@@ -161,10 +161,10 @@ def device(self) -> torch.device:
     def batch_size(self) -> int:
         return self._batch_size
 
-    def to(self, device: torch.device) -> None:
+    def to(self, device: Union[torch.device, str]) -> None:
         self._vertices = self._vertices.to(device=device)
         self._faces = self._faces.to(device=device)
-        self._device = device
+        self._device = torch.device(device) if isinstance(device, str) else device
 
 
 class Camera:
@@ -184,7 +184,7 @@ def __init__(
         resolution: Tuple[int, int],
         intrinsics: Optional[Union[torch.FloatTensor, np.ndarray]] = None,
         extrinsics: Optional[Union[torch.FloatTensor, np.ndarray]] = None,
-        device: torch.device = torch.device("cuda"),
+        device: Union[torch.device, str] = "cuda",
         name: str = "camera",
     ) -> None:
         if intrinsics is None:
@@ -211,7 +211,7 @@ def __init__(
         self._intrinsics = intrinsics
         self._extrinsics = extrinsics
         self._resolution = resolution
-        self._device = device
+        self._device = torch.device(device) if isinstance(device, str) else device
         ht_optical_shape = (
             (1,) + extrinsics.shape[-2:]
             if extrinsics.dim() == 3
@@ -229,11 +229,11 @@ def __init__(
         self._name = name
 
     @abc.abstractmethod
-    def to(self, device: torch.device) -> None:
+    def to(self, device: Union[torch.device, str]) -> None:
         self._intrinsics = self._intrinsics.to(device=device)
         self._extrinsics = self._extrinsics.to(device=device)
         self._ht_optical = self._ht_optical.to(device=device)
-        self._device = device
+        self._device = torch.device(device) if isinstance(device, str) else device
 
     @property
     def intrinsics(self) -> torch.FloatTensor:
@@ -312,7 +312,7 @@ def __init__(
         extrinsics: Optional[Union[torch.FloatTensor, np.ndarray]] = None,
         zmin: float = 0.1,
         zmax: float = 100.0,
-        device: torch.device = torch.device("cuda"),
+        device: Union[torch.device, str] = "cuda",
     ) -> None:
         super().__init__(resolution, intrinsics, extrinsics, device)
 
@@ -347,7 +347,7 @@ def __init__(
         self._perspective_projection[..., 2, 3] = 2.0 * zmax * zmin / (zmin - zmax)
         self._perspective_projection[..., 3, 2] = 1.0
 
-    def to(self, device: torch.device) -> None:
+    def to(self, device: Union[torch.device, str]) -> None:
         self._perspective_projection = self._perspective_projection.to(device=device)
         super().to(device=device)
 

roboreg/io/parsers.py

@@ -18,7 +18,7 @@ def __init__(self, urdf: str) -> None:
 
     @classmethod
     def from_file(cls, path: Union[Path, str]) -> None:
-        r"""Instantiate URDF parser from URDF string.
+        r"""Instantiate URDF parser path to URDF file.
 
         Args:
             path (Union[Path, str]): Path to URDF file.

roboreg/segmentor.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Union
 
 import numpy as np
 import torch
@@ -8,9 +8,11 @@
 class Segmentor(object):
     __slots__ = ["_model", "_pth", "_device"]
 
-    def __init__(self, pth: float = 0.5, device: str = "cuda") -> None:
+    def __init__(
+        self, pth: float = 0.5, device: Union[torch.device, str] = "cuda"
+    ) -> None:
         self._pth = pth
-        self._device = device
+        self._device = torch.device(device) if isinstance(device, str) else device
 
     @property
     def pth(self) -> float:
@@ -28,7 +30,7 @@ def __init__(
         self,
         model_id: str = "facebook/sam2-hiera-large",
         pth: float = 0.5,
-        device: str = "cuda",
+        device: Union[torch.device, str] = "cuda",
     ) -> None:
         super().__init__(pth=pth, device=device)
         self._model: SAM2ImagePredictor = SAM2ImagePredictor.from_pretrained(model_id)
@@ -37,7 +39,10 @@ def __call__(
         self, img: np.ndarray, input_points: np.ndarray, input_labels: np.ndarray
     ) -> np.ndarray:
         self._model.set_image(img)
-        with torch.inference_mode(), torch.autocast(self._device, dtype=torch.bfloat16):
+        with (
+            torch.inference_mode(),
+            torch.autocast(device_type=self._device.type, dtype=torch.bfloat16),
+        ):
             mask_logits, _, _ = self._model.predict(
                 point_coords=input_points,
                 point_labels=input_labels,

test/differentiable/test_structs.py

@@ -21,12 +21,14 @@ def test_torch_mesh_container() -> None:
     assert container.vertices.size()[1] == sum(
         list(container.per_mesh_vertex_count.values())
     ), "Expected vertex count to match."
-    assert container.device == device, f"Expected container on '{device}' device."
-    assert (
-        container.vertices[list(paths.keys())[0]].device == device
+    assert container.device == torch.device(
+        device
+    ), f"Expected container on '{device}' device."
+    assert container.vertices.device == torch.device(
+        device
     ), f"Expected vertices on '{device}' device."
-    assert (
-        container.faces[list(paths.keys())[0]].device == device
+    assert container.faces.device == torch.device(
+        device
     ), f"Expected faces on '{device}' device."