diff --git a/docs/en/understand_mmcv/ops.md b/docs/en/understand_mmcv/ops.md index ac9b36a5be..6c5f760bbe 100644 --- a/docs/en/understand_mmcv/ops.md +++ b/docs/en/understand_mmcv/ops.md @@ -46,7 +46,7 @@ We implement common ops used in detection, segmentation, etc. | RoIPool | | √ | √ | | √ | | RoIAlignRotated | √ | √ | √ | | | | RiRoIAlignRotated | | √ | | | | -| RoIAlign | √ | √ | √ | | | +| RoIAlign | √ | √ | √ | | √ | | RoIAwarePool3d | | √ | √ | | | | SAConv2d | | √ | | | | | SigmoidFocalLoss | | √ | √ | | √ | diff --git a/docs/zh_cn/understand_mmcv/ops.md b/docs/zh_cn/understand_mmcv/ops.md index 70179d05d2..0fd0873a32 100644 --- a/docs/zh_cn/understand_mmcv/ops.md +++ b/docs/zh_cn/understand_mmcv/ops.md @@ -46,7 +46,7 @@ MMCV 提供了检测、分割等任务中常用的算子 | RoIPool | | √ | √ | | √ | | RoIAlignRotated | √ | √ | √ | | | | RiRoIAlignRotated | | √ | | | | -| RoIAlign | √ | √ | √ | | | +| RoIAlign | √ | √ | √ | | √ | | RoIAwarePool3d | | √ | √ | | | | SAConv2d | | √ | | | | | SigmoidFocalLoss | | √ | √ | | √ | diff --git a/mmcv/ops/csrc/pytorch/npu/roi_align_npu.cpp b/mmcv/ops/csrc/pytorch/npu/roi_align_npu.cpp new file mode 100644 index 0000000000..31335557e4 --- /dev/null +++ b/mmcv/ops/csrc/pytorch/npu/roi_align_npu.cpp @@ -0,0 +1,68 @@ +#include "pytorch_npu_helper.hpp" + +using namespace NPU_NAME_SPACE; +using namespace std; + +void roi_align_forward_npu(Tensor input, Tensor rois, Tensor output, + Tensor argmax_y, Tensor argmax_x, int aligned_height, + int aligned_width, float spatial_scale, + int sampling_ratio, int pool_mode, bool aligned) { + if (!aligned) { + LOG(WARNING) << "The [aligned] attr in roi_align op is false"; + } + int64_t aligned_height_64 = aligned_height; + int64_t aligned_width_64 = aligned_width; + int64_t sampling_ratio_64 = sampling_ratio; + int64_t roi_end_mode = 0; + OpCommand cmd; + cmd.Name("ROIAlign") + .Input(input) + .Input(rois) + .Output(output) + .Attr("spatial_scale", spatial_scale) + .Attr("pooled_height", aligned_height_64) + .Attr("pooled_width", aligned_width_64) + .Attr("sample_num", sampling_ratio_64) + .Attr("roi_end_mode", roi_end_mode) + .Run(); +} + +void roi_align_backward_npu(Tensor grad_output, Tensor rois, Tensor argmax_y, + Tensor argmax_x, Tensor grad_input, + int aligned_height, int aligned_width, + float spatial_scale, int sampling_ratio, + int pool_mode, bool aligned) { + int64_t aligned_height_64 = aligned_height; + int64_t aligned_width_64 = aligned_width; + int64_t sampling_ratio_64 = sampling_ratio; + int64_t roi_end_mode = 0; + c10::SmallVector xdiff_shape = + at_npu::native::array_to_small_vector(grad_input.sizes()); + OpCommand cmd; + cmd.Name("ROIAlignGrad") + .Input(grad_output) + .Input(rois) + .Output(grad_input) + .Attr("xdiff_shape", xdiff_shape) + .Attr("pooled_width", aligned_width_64) + .Attr("pooled_height", aligned_height_64) + .Attr("spatial_scale", spatial_scale) + .Attr("sample_num", sampling_ratio_64) + .Attr("roi_end_mode", roi_end_mode) + .Run(); +} + +void roi_align_forward_impl(Tensor input, Tensor rois, Tensor output, + Tensor argmax_y, Tensor argmax_x, + int aligned_height, int aligned_width, + float spatial_scale, int sampling_ratio, + int pool_mode, bool aligned); + +void roi_align_backward_impl(Tensor grad_output, Tensor rois, Tensor argmax_y, + Tensor argmax_x, Tensor grad_input, + int aligned_height, int aligned_width, + float spatial_scale, int sampling_ratio, + int pool_mode, bool aligned); + +REGISTER_NPU_IMPL(roi_align_forward_impl, roi_align_forward_npu); +REGISTER_NPU_IMPL(roi_align_backward_impl, roi_align_backward_npu); diff --git a/tests/test_ops/test_roi_align.py b/tests/test_ops/test_roi_align.py index 6caf5c5356..9609f37c93 100644 --- a/tests/test_ops/test_roi_align.py +++ b/tests/test_ops/test_roi_align.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 _USING_PARROTS = True try: @@ -102,15 +102,19 @@ def _test_roialign_allclose(device, dtype): 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')) ]) @pytest.mark.parametrize('dtype', [ torch.float, pytest.param( torch.double, marks=pytest.mark.skipif( - IS_MLU_AVAILABLE, - reason='MLU does not support for 64-bit floating point')), + IS_MLU_AVAILABLE or IS_NPU_AVAILABLE, + reason='MLU and NPU do not support for 64-bit floating point')), torch.half ]) def test_roialign(device, dtype):