support musa backend in mmdetection3d

mmdet3d/apis/inference.py

@@ -13,6 +13,7 @@
 from mmengine.dataset import Compose, pseudo_collate
 from mmengine.registry import init_default_scope
 from mmengine.runner import load_checkpoint
+from mmengine.device import is_musa_available
 
 from mmdet3d.registry import DATASETS, MODELS
 from mmdet3d.structures import Box3DMode, Det3DDataSample, get_box_type
@@ -104,7 +105,10 @@ def init_model(config: Union[str, Path, Config],
 
     model.cfg = config  # save the config in the model for convenience
     if device != 'cpu':
-        torch.cuda.set_device(device)
+        if is_musa_available():
+            torch.musa.set_device(device)
+        else:
+            torch.cuda.set_device(device)
     else:
         warnings.warn('Don\'t suggest using CPU device. '
                       'Some functions are not supported for now.')

mmdet3d/evaluation/metrics/nuscenes_metric.py

@@ -10,6 +10,7 @@
 from mmengine import Config, load
 from mmengine.evaluator import BaseMetric
 from mmengine.logging import MMLogger
+from mmengine.device import is_musa_available
 from nuscenes.eval.detection.config import config_factory
 from nuscenes.eval.detection.data_classes import DetectionConfig
 from nuscenes.utils.data_classes import Box as NuScenesBox
@@ -777,11 +778,19 @@ def nusc_box_to_cam_box3d(
     dims[:, [0, 1, 2]] = dims[:, [1, 2, 0]]
     rots = -rots
 
-    boxes_3d = torch.cat([locs, dims, rots, velocity], dim=1).cuda()
-    cam_boxes3d = CameraInstance3DBoxes(
-        boxes_3d, box_dim=9, origin=(0.5, 0.5, 0.5))
-    scores = torch.Tensor([b.score for b in boxes]).cuda()
-    labels = torch.LongTensor([b.label for b in boxes]).cuda()
+    if is_musa_available():
+        boxes_3d = torch.cat([locs, dims, rots, velocity], dim=1).musa()
+        cam_boxes3d = CameraInstance3DBoxes(
+            boxes_3d, box_dim=9, origin=(0.5, 0.5, 0.5))
+        scores = torch.Tensor([b.score for b in boxes]).musa()
+        labels = torch.LongTensor([b.label for b in boxes]).musa()
+    else:
+        boxes_3d = torch.cat([locs, dims, rots, velocity], dim=1).cuda()
+        cam_boxes3d = CameraInstance3DBoxes(
+            boxes_3d, box_dim=9, origin=(0.5, 0.5, 0.5))
+        scores = torch.Tensor([b.score for b in boxes]).cuda()
+        labels = torch.LongTensor([b.label for b in boxes]).cuda()
+
     nms_scores = scores.new_zeros(scores.shape[0], 10 + 1)
     indices = labels.new_tensor(list(range(scores.shape[0])))
     nms_scores[indices, labels] = scores

mmdet3d/models/data_preprocessors/data_preprocessor.py

@@ -9,6 +9,7 @@
 from mmdet.models.utils.misc import samplelist_boxtype2tensor
 from mmengine.model import stack_batch
 from mmengine.utils import is_seq_of
+from mmengine.device import is_musa_available
 from torch import Tensor
 from torch.nn import functional as F
 
@@ -434,12 +435,18 @@ def voxelize(self, points: List[Tensor],
                 res_coors_numpy = res_coors.cpu().numpy()
                 inds, point2voxel_map = self.sparse_quantize(
                     res_coors_numpy, return_index=True, return_inverse=True)
-                point2voxel_map = torch.from_numpy(point2voxel_map).cuda()
+                if is_musa_available():
+                    point2voxel_map = torch.from_numpy(point2voxel_map).musa()
+                else:
+                    point2voxel_map = torch.from_numpy(point2voxel_map).cuda()
                 if self.training and self.max_voxels is not None:
                     if len(inds) > self.max_voxels:
                         inds = np.random.choice(
                             inds, self.max_voxels, replace=False)
-                inds = torch.from_numpy(inds).cuda()
+                if is_musa_available():
+                    inds = torch.from_numpy(inds).musa()
+                else:
+                    inds = torch.from_numpy(inds).cuda()
                 if hasattr(data_sample.gt_pts_seg, 'pts_semantic_mask'):
                     data_sample.gt_pts_seg.voxel_semantic_mask \
                         = data_sample.gt_pts_seg.pts_semantic_mask[inds]

mmdet3d/structures/bbox_3d/base_box3d.py

@@ -457,6 +457,18 @@ def cuda(self, *args, **kwargs) -> 'BaseInstance3DBoxes':
             box_dim=self.box_dim,
             with_yaw=self.with_yaw)
 
+    def musa(self, *args, **kwargs) -> 'BaseInstance3DBoxes':
+        """Convert current boxes to cuda device.
+
+        Returns:
+            :obj:`BaseInstance3DBoxes`: A new boxes object on the cuda device.
+        """
+        original_type = type(self)
+        return original_type(
+            self.tensor.musa(*args, **kwargs),
+            box_dim=self.box_dim,
+            with_yaw=self.with_yaw)
+
     def clone(self) -> 'BaseInstance3DBoxes':
         """Clone the boxes.
 

mmdet3d/structures/point_data.py

@@ -5,11 +5,21 @@
 import numpy as np
 import torch
 from mmengine.structures import BaseDataElement
+from mmengine.device import is_musa_available
 
 IndexType = Union[str, slice, int, list, torch.LongTensor,
                   torch.cuda.LongTensor, torch.BoolTensor,
                   torch.cuda.BoolTensor, np.ndarray]
 
+TensorType = (torch.BoolTensor, torch.cuda.BoolTensor)
+
+if is_musa_available():
+    IndexType = Union[str, slice, int, list, torch.LongTensor,
+                    torch.cuda.LongTensor, torch.BoolTensor,
+                    torch.cuda.BoolTensor, np.ndarray, torch.musa.LongTensor, torch.musa.BoolTensor]
+
+    TensorType = (torch.BoolTensor, torch.cuda.BoolTensor, torch.musa.BoolTensor)
+
 
 class PointData(BaseDataElement):
     """Data structure for point-level annotations or predictions.
@@ -85,9 +95,7 @@ def __getitem__(self, item: IndexType) -> 'PointData':
             # Mode details in https://github.com/numpy/numpy/issues/9464
             item = item.astype(np.int64) if item.dtype == np.int32 else item
             item = torch.from_numpy(item)
-        assert isinstance(
-            item, (str, slice, int, torch.LongTensor, torch.cuda.LongTensor,
-                   torch.BoolTensor, torch.cuda.BoolTensor))
+        assert isinstance(item, IndexType)
 
         if isinstance(item, str):
             return getattr(self, item)
@@ -103,7 +111,7 @@ def __getitem__(self, item: IndexType) -> 'PointData':
         if isinstance(item, torch.Tensor):
             assert item.dim() == 1, 'Only support to get the' \
                                     ' values along the first dimension.'
-            if isinstance(item, (torch.BoolTensor, torch.cuda.BoolTensor)):
+            if isinstance(item, TensorType):
                 assert len(item) == len(self), 'The shape of the ' \
                                                'input(BoolTensor) ' \
                                                f'{len(item)} ' \
@@ -122,8 +130,7 @@ def __getitem__(self, item: IndexType) -> 'PointData':
                         v, (str, list, tuple)) or (hasattr(v, '__getitem__')
                                                    and hasattr(v, 'cat')):
                     # convert to indexes from BoolTensor
-                    if isinstance(item,
-                                  (torch.BoolTensor, torch.cuda.BoolTensor)):
+                    if isinstance(item, TensorType):
                         indexes = torch.nonzero(item).view(
                             -1).cpu().numpy().tolist()
                     else:

mmdet3d/structures/points/base_points.py

@@ -461,6 +461,18 @@ def cuda(self, *args, **kwargs) -> 'BasePoints':
             points_dim=self.points_dim,
             attribute_dims=self.attribute_dims)
 
+    def musa(self, *args, **kwargs) -> 'BasePoints':
+        """Convert current points to cuda device.
+
+        Returns:
+            :obj:`BasePoints`: A new points object on the cuda device.
+        """
+        original_type = type(self)
+        return original_type(
+            self.tensor.musa(*args, **kwargs),
+            points_dim=self.points_dim,
+            attribute_dims=self.attribute_dims)
+
     def clone(self) -> 'BasePoints':
         """Clone the points.
 

mmdet3d/testing/model_utils.py

@@ -6,18 +6,34 @@
 import numpy as np
 import torch
 from mmengine.structures import InstanceData
+from mmengine.device import is_musa_available, is_cuda_available
 
 from mmdet3d.structures import (CameraInstance3DBoxes, DepthInstance3DBoxes,
                                 Det3DDataSample, LiDARInstance3DBoxes,
                                 PointData)
 
 
+AVAILABLE_DEVICES = []
+
+if is_cuda_available():
+    AVAILABLE_DEVICES.append(("cuda", "cuda"))
+
+if is_musa_available():
+    AVAILABLE_DEVICES.append(("musa", "musa"))
+
+
 def setup_seed(seed):
     torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
+    if is_musa_available():
+        torch.musa.manual_seed_all(seed)
+    else:
+        torch.cuda.manual_seed_all(seed)
     np.random.seed(seed)
     random.seed(seed)
-    torch.backends.cudnn.deterministic = True
+    if is_musa_available():
+        torch.backends.mudnn.deterministic = True
+    else:
+        torch.backends.cudnn.deterministic = True
 
 
 def _get_config_directory():

tests/test_apis/test_inferencers/test_lidar_det3d_inferencer.py

@@ -7,6 +7,7 @@
 import numpy as np
 import torch
 from mmengine.utils import is_list_of
+from mmengine.device import is_cuda_available, is_musa_available
 
 from mmdet3d.apis import LidarDet3DInferencer
 from mmdet3d.structures import Det3DDataSample
@@ -41,7 +42,7 @@ def assert_predictions_equal(self, preds1, preds2):
                     np.allclose(pred1['labels_3d'], pred2['labels_3d']))
 
     def test_call(self):
-        if not torch.cuda.is_available():
+        if not is_cuda_available() and not is_musa_available():
             return
         # single point cloud
         inputs = dict(points='tests/data/kitti/training/velodyne/000000.bin')
@@ -84,7 +85,7 @@ def test_call(self):
         self.assertIn('predictions', res_bs2)
 
     def test_visualize(self):
-        if not torch.cuda.is_available():
+        if not is_cuda_available() and not is_musa_available():
             return
         inputs = dict(points='tests/data/kitti/training/velodyne/000000.bin'),
         # img_out_dir
@@ -95,7 +96,7 @@ def test_visualize(self):
             # self.assertTrue(osp.exists(osp.join(tmp_dir, '000000.png')))
 
     def test_postprocess(self):
-        if not torch.cuda.is_available():
+        if not is_cuda_available() and not is_musa_available():
             return
         # return_datasample
         inputs = dict(points='tests/data/kitti/training/velodyne/000000.bin')

tests/test_apis/test_inferencers/test_lidar_seg3d_inferencer.py

@@ -9,6 +9,7 @@
 import pytest
 import torch
 from mmengine.utils import is_list_of
+from mmengine.device import is_cuda_available, is_musa_available
 
 from mmdet3d.apis import LidarSeg3DInferencer
 from mmdet3d.structures import Det3DDataSample
@@ -36,7 +37,9 @@ def assert_predictions_equal(self, preds1, preds2):
                             pred2['pts_semantic_mask']))
 
     @pytest.mark.skipif(
-        not torch.cuda.is_available(), reason='requires CUDA support')
+        not is_cuda_available(), reason='requires CUDA support')
+    @pytest.mark.skipif(
+        not is_musa_available(), reason='requires MUSA support')
     @pytest.mark.skipif(
         'DISPLAY' not in os.environ, reason='requires DISPLAY device')
     def test_call(self):
@@ -84,7 +87,9 @@ def test_call(self):
         self.assertIn('predictions', res_bs2)
 
     @pytest.mark.skipif(
-        not torch.cuda.is_available(), reason='requires CUDA support')
+        not is_cuda_available(), reason='requires CUDA support')
+    @pytest.mark.skipif(
+        not is_musa_available(), reason='requires MUSA support')
     @pytest.mark.skipif(
         'DISPLAY' not in os.environ, reason='requires DISPLAY device')
     def test_visualizer(self):
@@ -94,7 +99,7 @@ def test_visualizer(self):
             self.inferencer(inputs, out_dir=tmp_dir)
 
     def test_post_processor(self):
-        if not torch.cuda.is_available():
+        if not is_cuda_available() and not is_musa_available():
             return
         # return_datasample
         inputs = dict(points='tests/data/s3dis/points/Area_1_office_2.bin')

tests/test_apis/test_inferencers/test_multi_modality_det3d_inferencer.py

@@ -8,6 +8,7 @@
 import numpy as np
 import torch
 from mmengine.utils import is_list_of
+from mmengine.device import is_cuda_available, is_musa_available
 
 from mmdet3d.apis import MultiModalityDet3DInferencer
 from mmdet3d.structures import Det3DDataSample
@@ -42,7 +43,7 @@ def assert_predictions_equal(self, preds1, preds2):
                     np.allclose(pred1['labels_3d'], pred2['labels_3d']))
 
     def test_call(self):
-        if not torch.cuda.is_available():
+        if not is_cuda_available() and not is_musa_available():
             return
         infos_path = 'demo/data/kitti/000008.pkl'
         points_path = 'demo/data/kitti/000008.bin'
@@ -86,7 +87,7 @@ def test_call(self):
         self.assertIn('visualization', res_ndarray)
 
     def test_visualize(self):
-        if not torch.cuda.is_available():
+        if not is_cuda_available() and not is_musa_available():
             return
         inputs = dict(
             points='demo/data/kitti/000008.bin',
@@ -100,7 +101,7 @@ def test_visualize(self):
             # self.assertTrue(osp.exists(osp.join(tmp_dir, '000000.png')))
 
     def test_postprocess(self):
-        if not torch.cuda.is_available():
+        if not is_cuda_available() and not is_musa_available():
             return
         # return_datasample
         infos_path = 'demo/data/kitti/000008.pkl'

tests/test_evaluation/test_functional/test_kitti_eval.py

@@ -2,13 +2,14 @@
 import numpy as np
 import pytest
 import torch
+from mmengine.device import is_cuda_available, is_musa_available
 
 from mmdet3d.evaluation import do_eval, eval_class, kitti_eval
 
 
 def test_do_eval():
-    if not torch.cuda.is_available():
-        pytest.skip('test requires GPU and CUDA')
+    if not is_cuda_available() and not is_musa_available():
+        pytest.skip('test requires GPU and CUDA or MUSA')
     gt_name = np.array(
         ['Pedestrian', 'Cyclist', 'Car', 'Car', 'Car', 'DontCare', 'DontCare'])
     gt_truncated = np.array([0., 0., 0., -1., -1., -1., -1.])
@@ -128,8 +129,8 @@ def test_do_eval():
 
 
 def test_kitti_eval():
-    if not torch.cuda.is_available():
-        pytest.skip('test requires GPU and CUDA')
+    if not is_cuda_available() and not is_musa_available():
+        pytest.skip('test requires GPU and CUDA or MUSA')
     gt_name = np.array(
         ['Pedestrian', 'Cyclist', 'Car', 'Car', 'Car', 'DontCare', 'DontCare'])
     gt_truncated = np.array([0., 0., 0., -1., -1., -1., -1.])

tests/test_evaluation/test_functional/test_panoptic_seg_eval.py

@@ -2,12 +2,13 @@
 import numpy as np
 import pytest
 import torch
+from mmengine.device import is_cuda_available, is_musa_available
 
 from mmdet3d.evaluation.functional.panoptic_seg_eval import panoptic_seg_eval
 
 
 def test_panoptic_seg_eval():
-    if not torch.cuda.is_available():
+    if not is_cuda_available() and not is_musa_available():
         pytest.skip()
 
     classes = ['unlabeled', 'person', 'dog', 'grass', 'sky']

tests/test_evaluation/test_functional/test_seg_eval.py

@@ -2,12 +2,12 @@
 import numpy as np
 import pytest
 import torch
-
+from mmengine.device import is_cuda_available, is_musa_available
 from mmdet3d.evaluation.functional.seg_eval import seg_eval
 
 
 def test_indoor_eval():
-    if not torch.cuda.is_available():
+    if not is_cuda_available() and not is_musa_available():
         pytest.skip()
     seg_preds = [
         np.array([