diff --git a/mmcv/ops/csrc/pytorch/npu/roi_align_rotated_v2_npu.cpp b/mmcv/ops/csrc/pytorch/npu/roi_align_rotated_v2_npu.cpp index b2ea93b261..64248ada45 100644 --- a/mmcv/ops/csrc/pytorch/npu/roi_align_rotated_v2_npu.cpp +++ b/mmcv/ops/csrc/pytorch/npu/roi_align_rotated_v2_npu.cpp @@ -3,25 +3,36 @@ using namespace NPU_NAME_SPACE; using namespace std; -void roi_align_rotated_v2_forward_npu(const Tensor input, Tensor rois_map, - Tensor output, +void roi_align_rotated_v2_forward_npu(const Tensor x, Tensor rois_map, + Tensor y, + int32_t pooled_h, + int32_t pooled_w, double spatial_scale, int32_t sampling_ratio, - int32_t pooled_height, - int32_t pooled_width, bool aligned, bool clockwise) { - at::Tensor feature_map = input.permute({0, 2, 3, 1}).contiguous(); + at::Tensor feature_map = x.permute({0, 2, 3, 1}).contiguous(); at::Tensor rois = rois_map.permute({1, 0}).contiguous(); - EXEC_NPU_CMD(aclnnRoiAlignRotatedV2, feature_map, rois, spatial_scale, sampling_ratio, pooled_height, pooled_width, aligned, clockwise, output); + at_npu::native::OpCommand cmd; + cmd.Name("RoiAlignRotated") + .Input(feature_map) + .Input(rois) + .Output(y) + .Attr("pooled_h", static_cast(pooled_h)) + .Attr("pooled_w", static_cast(pooled_w)) + .Attr("spatial_scale", static_cast(spatial_scale)) + .Attr("sampling_ratio", static_cast(sampling_ratio)) + .Attr("aligned", aligned) + .Attr("clockwise", clockwise) + .Run(); } -void roi_align_rotated_v2_forward_impl(const Tensor input, Tensor rois, - Tensor output, +void roi_align_rotated_v2_forward_impl(const Tensor x, Tensor rois, + Tensor y, + int32_t pooled_h, + int32_t pooled_w, double spatial_scale, int32_t sampling_ratio, - int32_t pooled_height, - int32_t pooled_width, bool aligned, bool clockwise); diff --git a/mmcv/ops/csrc/pytorch/pybind.cpp b/mmcv/ops/csrc/pytorch/pybind.cpp index 3745bd8978..43e9b270f2 100644 --- a/mmcv/ops/csrc/pytorch/pybind.cpp +++ b/mmcv/ops/csrc/pytorch/pybind.cpp @@ -209,8 +209,8 @@ void roi_align_backward(Tensor grad_output, Tensor rois, Tensor argmax_y, int sampling_ratio, int pool_mode, bool aligned); void roi_align_rotated_v2_forward(Tensor input, Tensor rois, Tensor output, + int pooled_h, int pooled_w, double spatial_scale, int sampling_ratio, - int aligned_height, int aligned_width, bool aligned, bool clockwise); void roi_align_rotated_v2_backward(Tensor input, Tensor rois, @@ -343,9 +343,9 @@ void roi_align_rotated_backward(Tensor grad_output, Tensor rois, bool clockwise); void roi_align_rotated_v2_forward(Tensor input, Tensor rois, Tensor output, - double spatial_scale, int sampling_ratio, - int aligned_height, int aligned_width, - bool aligned, bool clockwise); + int pooled_h, int pooled_w, + double spatial_scale, int sampling_ratio, + bool aligned, bool clockwise); void roi_align_rotated_v2_backward(Tensor input, Tensor rois, Tensor grad_output, Tensor grad_input, int pooled_height, int pooled_width, double spatial_scale, @@ -814,8 +814,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("sampling_ratio"), py::arg("aligned"), py::arg("clockwise")); m.def("roi_align_rotated_v2_forward", &roi_align_rotated_v2_forward, "roi_align_rotated_v2_forward", py::arg("input"), py::arg("rois"), - py::arg("output"), py::arg("spatial_scale"), py::arg("sampling_ratio"), - py::arg("pooled_height"), py::arg("pooled_width"), + py::arg("output"), py::arg("pooled_h"), py::arg("pooled_w"), + py::arg("spatial_scale"), py::arg("sampling_ratio"), py::arg("aligned"), py::arg("clockwise")); m.def("roi_align_rotated_v2_backward", &roi_align_rotated_v2_backward, "roi_align_rotated_v2_backward", py::arg("input"), py::arg("rois"), diff --git a/mmcv/ops/csrc/pytorch/roi_align_rotated_v2.cpp b/mmcv/ops/csrc/pytorch/roi_align_rotated_v2.cpp index 7743775288..ec2a17bf77 100644 --- a/mmcv/ops/csrc/pytorch/roi_align_rotated_v2.cpp +++ b/mmcv/ops/csrc/pytorch/roi_align_rotated_v2.cpp @@ -2,22 +2,22 @@ #include "pytorch_cpp_helper.hpp" #include "pytorch_device_registry.hpp" -void roi_align_rotated_v2_forward_impl(Tensor input, Tensor rois, Tensor output, +void roi_align_rotated_v2_forward_impl(Tensor x, Tensor rois, Tensor y, + int pooled_h, int pooled_w, double spatial_scale, int sampling_ratio, - int pooled_height, int pooled_width, bool aligned, bool clockwise) { - DISPATCH_DEVICE_IMPL(roi_align_rotated_v2_forward_impl, input, rois, output, - spatial_scale, sampling_ratio, pooled_height, pooled_width, + DISPATCH_DEVICE_IMPL(roi_align_rotated_v2_forward_impl, x, rois, y, + pooled_h, pooled_w, spatial_scale, sampling_ratio, aligned, clockwise); } -void roi_align_rotated_v2_forward(Tensor input, Tensor rois, Tensor output, +void roi_align_rotated_v2_forward(Tensor x, Tensor rois, Tensor y, + int pooled_h, int pooled_w, double spatial_scale, int sampling_ratio, - int pooled_height, int pooled_width, bool aligned, bool clockwise) { - roi_align_rotated_v2_forward_impl(input, rois, output, spatial_scale, sampling_ratio, - pooled_height, pooled_width, aligned, clockwise); + roi_align_rotated_v2_forward_impl(x, rois, y, pooled_h, pooled_w, + spatial_scale, sampling_ratio, aligned, clockwise); } diff --git a/mmcv/ops/roi_align_rotated_v2.py b/mmcv/ops/roi_align_rotated_v2.py index 639fea3a23..80c97ae736 100644 --- a/mmcv/ops/roi_align_rotated_v2.py +++ b/mmcv/ops/roi_align_rotated_v2.py @@ -14,55 +14,55 @@ class RoIAlignRotatedV2Function(Function): @staticmethod - def symbolic(g, input, rois, spatial_scale, sampling_ratio, pooled_height, - pooled_width, aligned, clockwise): + def symbolic(g, x, rois, spatial_scale, sampling_ratio, pooled_h, + pooled_w, aligned, clockwise): return g.op( 'mmcv::MMCVRoIAlignRotatedV2', - input, + x, rois, + pooled_h=pooled_h, + pooled_w=pooled_w, spatial_scale_f=spatial_scale, sampling_ratio_i=sampling_ratio, - pooled_height=pooled_height, - pooled_width=pooled_width, aligned_i=aligned, clockwise_i=clockwise) @staticmethod def forward(ctx: Any, - input: torch.Tensor, + x: torch.Tensor, rois: torch.Tensor, + pooled_h: int, + pooled_w: int, spatial_scale: float, sampling_ratio: int, - pooled_height: int, - pooled_width: int, aligned: bool = True, clockwise: bool = False) -> torch.Tensor: - ctx.pooled_height = pooled_height - ctx.pooled_width = pooled_width + ctx.pooled_h = pooled_h + ctx.pooled_w = pooled_w ctx.spatial_scale = spatial_scale ctx.sampling_ratio = sampling_ratio ctx.aligned = aligned ctx.clockwise = clockwise - ctx.save_for_backward(input, rois) - ctx.feature_size = input.size() - batch_size, num_channels, data_height, data_width = input.size() + ctx.save_for_backward(x, rois) + ctx.feature_size = x.size() + batch_size, num_channels, data_height, data_width = x.size() num_rois = rois.size(0) - output = input.new_zeros(num_rois, ctx.pooled_height, ctx.pooled_width, + y = x.new_zeros(num_rois, ctx.pooled_h, ctx.pooled_w, num_channels) ext_module.roi_align_rotated_v2_forward( - input, + x, rois, - output, + y, + pooled_h=ctx.pooled_h, + pooled_w=ctx.pooled_w, spatial_scale=ctx.spatial_scale, sampling_ratio=ctx.sampling_ratio, - pooled_height=ctx.pooled_height, - pooled_width=ctx.pooled_width, aligned=ctx.aligned, clockwise=ctx.clockwise) - output = output.transpose(2, 3).transpose(1, 2).contiguous() - return output + y = y.transpose(2, 3).transpose(1, 2).contiguous() + return y @staticmethod def backward(ctx: Any, grad_output: torch.Tensor): @@ -74,7 +74,7 @@ def backward(ctx: Any, grad_output: torch.Tensor): input.size(0), input.size(2), input.size(3), input.size(1)) ext_module.roi_align_rotated_v2_backward( input, rois_trans, grad_output_trans, grad_input, - ctx.pooled_height, ctx.pooled_width, ctx.spatial_scale, + ctx.pooled_h, ctx.pooled_w, ctx.spatial_scale, ctx.sampling_ratio, ctx.aligned, ctx.clockwise) grad_input = grad_input.permute(0, 3, 1, 2).contiguous() @@ -134,31 +134,33 @@ class RoIAlignRotatedV2(nn.Module): }, cls_name='RoIAlignRotatedV2') def __init__(self, + pooled_h: int, + pooled_w: int, spatial_scale: float, sampling_ratio: int, - pooled_height: int, - pooled_width: int, aligned: bool = True, clockwise: bool = False): super().__init__() - self.pooled_height = int(pooled_height) - self.pooled_width = int(pooled_width) + self.pooled_h = int(pooled_h) + self.pooled_w = int(pooled_w) self.spatial_scale = float(spatial_scale) self.sampling_ratio = int(sampling_ratio) self.aligned = aligned self.clockwise = clockwise def forward(self, input: torch.Tensor, rois: torch.Tensor) -> torch.Tensor: - return RoIAlignRotatedV2Function.apply(input, rois, self.spatial_scale, + return RoIAlignRotatedV2Function.apply(input, rois, + self.pooled_h, + self.pooled_w, + self.spatial_scale, self.sampling_ratio, - self.pooled_height, - self.pooled_width, self.aligned, + self.aligned, self.clockwise) def __repr__(self): s = self.__class__.__name__ - s += f'(pooled_height={self.pooled_height}, ' + s += f'(pooled_h={self.pooled_h}, ' s += f'spatial_scale={self.spatial_scale}, ' s += f'sampling_ratio={self.sampling_ratio}, ' s += f'aligned={self.aligned}, '