Merge pull request #2390 from zcc-zjut/mask

Add masked_ Conv2d operator in NPU
open-mmlab · Nov 7, 2022 · e3d78a8 · e3d78a8
2 parents 57a2734 + aa90e8e
commit e3d78a8
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/mmcv/ops/masked_conv.py b/mmcv/ops/masked_conv.py
@@ -45,6 +45,22 @@ def forward(ctx,
                 'Stride could not only be 1 in masked_conv2d currently.')
         out_channel, in_channel, kernel_h, kernel_w = weight.size()
 
+        if features.device.type == 'npu':
+            import torch_npu
+            conv = torch_npu.npu_conv2d(
+                features,
+                weight,
+                bias,
+                stride=(stride_h, stride_w),
+                padding=padding,
+                dilation=(1, 1),
+                groups=1)
+            features_h, features_w = features.size()[2:]
+            mask_reshape = mask.reshape(1, 1, features_h, features_w)
+            mask_bool = mask_reshape > 0
+            output = conv * mask_bool
+            return output
+
         batch_size = features.size(0)
         out_h = int(
             math.floor(

diff --git a/tests/test_ops/test_masked_conv2d.py b/tests/test_ops/test_masked_conv2d.py
@@ -3,7 +3,7 @@
 import pytest
 import torch
 
-from mmcv.utils import IS_CUDA_AVAILABLE, IS_MLU_AVAILABLE
+from mmcv.utils import IS_CUDA_AVAILABLE, IS_MLU_AVAILABLE, IS_NPU_AVAILABLE
 
 
 class TestMaskedConv2d:
@@ -16,7 +16,11 @@ class TestMaskedConv2d:
         pytest.param(
             'mlu',
             marks=pytest.mark.skipif(
-                not IS_MLU_AVAILABLE, reason='requires MLU support'))
+                not IS_MLU_AVAILABLE, reason='requires MLU support')),
+        pytest.param(
+            'npu',
+            marks=pytest.mark.skipif(
+                not IS_NPU_AVAILABLE, reason='requires NPU support'))
     ])
     def test_masked_conv2d_all_close(self, device):
         from mmcv.ops import MaskedConv2d