support moore threads MUSA gpu

mmseg/apis/mmseg_inferencer.py

@@ -9,6 +9,7 @@
 import torch
 import torch.nn as nn
 from mmcv.transforms import Compose
+from mmengine.device.utils import is_musa_available
 from mmengine.infer.infer import BaseInferencer, ModelType
 from mmengine.model import revert_sync_batchnorm
 from mmengine.registry import init_default_scope
@@ -81,7 +82,8 @@ def __init__(self,
         super().__init__(
             model=model, weights=weights, device=device, scope=scope)
 
-        if device == 'cpu' or not torch.cuda.is_available():
+        if device == 'cpu' or (not torch.cuda.is_available()
+                               and not is_musa_available()):
             self.model = revert_sync_batchnorm(self.model)
 
         assert isinstance(self.visualizer, SegLocalVisualizer)

mmseg/models/assigners/hungarian_assigner.py

@@ -3,9 +3,14 @@
 
 import torch
 from mmengine import ConfigDict
+from mmengine.device.utils import is_musa_available
 from mmengine.structures import InstanceData
 from scipy.optimize import linear_sum_assignment
-from torch.cuda.amp import autocast
+
+if is_musa_available():
+    from torch_musa.core.amp import autocast
+else:
+    from torch.cuda.amp import autocast
 
 from mmseg.registry import TASK_UTILS
 from .base_assigner import BaseAssigner

mmseg/models/losses/focal_loss.py

@@ -4,6 +4,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from mmcv.ops import sigmoid_focal_loss as _sigmoid_focal_loss
+from mmengine.device.utils import is_musa_available
 
 from mmseg.registry import MODELS
 from .utils import weight_reduce_loss
@@ -269,7 +270,9 @@ def forward(self,
             reduction_override if reduction_override else self.reduction)
         if self.use_sigmoid:
             num_classes = pred.size(1)
-            if torch.cuda.is_available() and pred.is_cuda:
+            if (torch.cuda.is_available()
+                    and pred.is_cuda) or (is_musa_available()
+                                          and pred.device == 'musa'):
                 if target.dim() == 1:
                     one_hot_target = F.one_hot(
                         target, num_classes=num_classes + 1)

tests/test_models/test_backbones/test_clip_text_encoder.py

@@ -1,6 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
 from mmengine import Config
+from mmengine.device.utils import is_musa_available
 from mmengine.registry import init_default_scope
 
 from mmseg.models.text_encoder import CLIPTextEncoder
@@ -23,6 +24,8 @@ def test_clip_text_encoder():
     text_encoder = CLIPTextEncoder(**cfg)
     if torch.cuda.is_available():
         text_encoder = text_encoder.cuda()
+    elif is_musa_available():
+        text_encoder = text_encoder.musa()
 
     with torch.no_grad():
         class_embeds = text_encoder()

tests/test_models/test_forward.py

@@ -8,6 +8,7 @@
 import pytest
 import torch
 import torch.nn as nn
+from mmengine.device.utils import is_musa_available
 from mmengine.model.utils import revert_sync_batchnorm
 from mmengine.registry import init_default_scope
 from mmengine.structures import PixelData
@@ -190,6 +191,8 @@ def _test_encoder_decoder_forward(cfg_file):
     # convert to cuda Tensor if applicable
     if torch.cuda.is_available():
         segmentor = segmentor.cuda()
+    elif is_musa_available():
+        segmentor = segmentor.musa()
     else:
         segmentor = revert_sync_batchnorm(segmentor)
 

tests/test_models/test_heads/test_ann_head.py

@@ -1,8 +1,9 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmengine.device.utils import is_musa_available
 
 from mmseg.models.decode_heads import ANNHead
-from .utils import to_cuda
+from .utils import to_cuda, to_musa
 
 
 def test_ann_head():
@@ -16,5 +17,7 @@ def test_ann_head():
         project_channels=8)
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     outputs = head(inputs)
     assert outputs.shape == (1, head.num_classes, 21, 21)

tests/test_models/test_heads/test_apc_head.py

@@ -1,9 +1,10 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import pytest
 import torch
+from mmengine.device.utils import is_musa_available
 
 from mmseg.models.decode_heads import APCHead
-from .utils import _conv_has_norm, to_cuda
+from .utils import _conv_has_norm, to_cuda, to_musa
 
 
 def test_apc_head():
@@ -34,6 +35,8 @@ def test_apc_head():
         fusion=True)
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     assert head.fusion is True
     assert head.acm_modules[0].pool_scale == 1
     assert head.acm_modules[1].pool_scale == 2
@@ -51,6 +54,8 @@ def test_apc_head():
         fusion=False)
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     assert head.fusion is False
     assert head.acm_modules[0].pool_scale == 1
     assert head.acm_modules[1].pool_scale == 2

tests/test_models/test_heads/test_aspp_head.py

@@ -1,9 +1,10 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import pytest
 import torch
+from mmengine.device.utils import is_musa_available
 
 from mmseg.models.decode_heads import ASPPHead, DepthwiseSeparableASPPHead
-from .utils import _conv_has_norm, to_cuda
+from .utils import _conv_has_norm, to_cuda, to_musa
 
 
 def test_aspp_head():
@@ -29,6 +30,8 @@ def test_aspp_head():
         in_channels=8, channels=4, num_classes=19, dilations=(1, 12, 24))
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     assert head.aspp_modules[0].conv.dilation == (1, 1)
     assert head.aspp_modules[1].conv.dilation == (12, 12)
     assert head.aspp_modules[2].conv.dilation == (24, 24)
@@ -49,6 +52,8 @@ def test_dw_aspp_head():
         dilations=(1, 12, 24))
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     assert head.c1_bottleneck is None
     assert head.aspp_modules[0].conv.dilation == (1, 1)
     assert head.aspp_modules[1].depthwise_conv.dilation == (12, 12)
@@ -67,6 +72,8 @@ def test_dw_aspp_head():
         dilations=(1, 12, 24))
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     assert head.c1_bottleneck.in_channels == 4
     assert head.c1_bottleneck.out_channels == 2
     assert head.aspp_modules[0].conv.dilation == (1, 1)

tests/test_models/test_heads/test_decode_head.py

@@ -3,11 +3,12 @@
 
 import pytest
 import torch
+from mmengine.device.utils import is_musa_available
 from mmengine.structures import PixelData
 
 from mmseg.models.decode_heads.decode_head import BaseDecodeHead
 from mmseg.structures import SegDataSample
-from .utils import to_cuda
+from .utils import to_cuda, to_musa
 
 
 @patch.multiple(BaseDecodeHead, __abstractmethods__=set())
@@ -70,6 +71,8 @@ def test_decode_head():
     head = BaseDecodeHead(32, 16, num_classes=19)
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     assert head.in_channels == 32
     assert head.input_transform is None
     transformed_inputs = head._transform_inputs(inputs)
@@ -84,6 +87,8 @@ def test_decode_head():
                           input_transform='resize_concat')
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     assert head.in_channels == 48
     assert head.input_transform == 'resize_concat'
     transformed_inputs = head._transform_inputs(inputs)
@@ -108,6 +113,8 @@ def test_decode_head():
             type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0))
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     loss = head.loss_by_feat(
         seg_logits=inputs, batch_data_samples=data_samples)
     assert 'loss_ce' in loss
@@ -128,6 +135,8 @@ def test_decode_head():
         ])
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
 
     loss = head.loss_by_feat(
         seg_logits=inputs, batch_data_samples=data_samples)
@@ -155,6 +164,8 @@ def test_decode_head():
                      dict(type='CrossEntropyLoss', loss_name='loss_3')))
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     loss = head.loss_by_feat(
         seg_logits=inputs, batch_data_samples=data_samples)
     assert 'loss_1' in loss
@@ -176,6 +187,8 @@ def test_decode_head():
                      dict(type='CrossEntropyLoss', loss_name='loss_ce')))
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     loss_3 = head.loss_by_feat(
         seg_logits=inputs, batch_data_samples=data_samples)
 
@@ -186,6 +199,8 @@ def test_decode_head():
         loss_decode=(dict(type='CrossEntropyLoss', loss_name='loss_ce')))
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     loss = head.loss_by_feat(
         seg_logits=inputs, batch_data_samples=data_samples)
     assert 'loss_ce' in loss

tests/test_models/test_heads/test_dm_head.py

@@ -1,9 +1,10 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import pytest
 import torch
+from mmengine.device.utils import is_musa_available
 
 from mmseg.models.decode_heads import DMHead
-from .utils import _conv_has_norm, to_cuda
+from .utils import _conv_has_norm, to_cuda, to_musa
 
 
 def test_dm_head():
@@ -34,6 +35,8 @@ def test_dm_head():
         fusion=True)
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     assert head.fusion is True
     assert head.dcm_modules[0].filter_size == 1
     assert head.dcm_modules[1].filter_size == 3
@@ -51,6 +54,8 @@ def test_dm_head():
         fusion=False)
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     assert head.fusion is False
     assert head.dcm_modules[0].filter_size == 1
     assert head.dcm_modules[1].filter_size == 3

tests/test_models/test_heads/test_dnl_head.py

@@ -1,8 +1,9 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmengine.device.utils import is_musa_available
 
 from mmseg.models.decode_heads import DNLHead
-from .utils import to_cuda
+from .utils import to_cuda, to_musa
 
 
 def test_dnl_head():
@@ -14,6 +15,8 @@ def test_dnl_head():
     inputs = [torch.randn(1, 8, 23, 23)]
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     outputs = head(inputs)
     assert outputs.shape == (1, head.num_classes, 23, 23)
 
@@ -23,6 +26,8 @@ def test_dnl_head():
     inputs = [torch.randn(1, 8, 23, 23)]
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     outputs = head(inputs)
     assert outputs.shape == (1, head.num_classes, 23, 23)
 
@@ -31,6 +36,8 @@ def test_dnl_head():
     inputs = [torch.randn(1, 8, 23, 23)]
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     outputs = head(inputs)
     assert outputs.shape == (1, head.num_classes, 23, 23)
 
@@ -40,5 +47,7 @@ def test_dnl_head():
     inputs = [torch.randn(1, 8, 23, 23)]
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     outputs = head(inputs)
     assert outputs.shape == (1, head.num_classes, 23, 23)

tests/test_models/test_heads/test_ema_head.py

@@ -1,8 +1,9 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
+from mmengine.device.utils import is_musa_available
 
 from mmseg.models.decode_heads import EMAHead
-from .utils import to_cuda
+from .utils import to_cuda, to_musa
 
 
 def test_emanet_head():
@@ -19,5 +20,7 @@ def test_emanet_head():
     inputs = [torch.randn(1, 4, 23, 23)]
     if torch.cuda.is_available():
         head, inputs = to_cuda(head, inputs)
+    elif is_musa_available():
+        head, inputs = to_musa(head, inputs)
     outputs = head(inputs)
     assert outputs.shape == (1, head.num_classes, 23, 23)