add roi_align op adpater

open-mmlab · Mar 4, 2023 · 95cd764 · 95cd764
1 parent 1c8af8e
commit 95cd764
Showing 1 changed file with 52 additions and 44 deletions.
diff --git a/mmcv/ops/csrc/pytorch/npu/roi_align_npu.cpp b/mmcv/ops/csrc/pytorch/npu/roi_align_npu.cpp
@@ -3,61 +3,69 @@
 using namespace NPU_NAME_SPACE;
 using namespace std;
 
-void roi_pool_forward_npu(Tensor input, Tensor rois, Tensor output,
-                          Tensor argmax, int pooled_height, int pooled_width,
-                          float spatial_scale) {
-  int64_t pooled_height_64 = pooled_height;
-  int64_t pooled_width_64 = pooled_width;
-  int64_t pooled_channel = 1;
-  at::Tensor roi_actual_num = at_npu::native::OpPreparation::ApplyTensor(
-      {}, rois.options().dtype(at::kInt), rois);
+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) {
+  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 = 1;
+  auto outputSize ={rois.size(0),input.size(1),aligned_height_64,aligned_width_64};
+
+  #at::Tensor result =
+  #    OpPreparation::ApplyTensorWithFormat(input, outputSize, ACL_FORMAT_NC1HWC0);
   OpCommand cmd;
-  cmd.Name("RoiPoolingWithArgMax")
+  cmd.Name("ROIAlign")
       .Input(input)
       .Input(rois)
-      .Input(roi_actual_num)
-      .Output(output)
-      .Output(argmax)
-      .Attr("pooled_h", pooled_height_64)
-      .Attr("pooled_w", pooled_width_64)
-      .Attr("spatial_scale_h", spatial_scale)
-      .Attr("spatial_scale_w", spatial_scale)
-      .Attr("pool_channel", pooled_channel)
+      .Output(result)
+      .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_pool_backward_npu(Tensor grad_output, Tensor rois, Tensor argmax,
-                           Tensor grad_input, int pooled_height,
-                           int pooled_width, float spatial_scale) {
-  int64_t pooled_height_64 = pooled_height;
-  int64_t pooled_width_64 = pooled_width;
-  int64_t pooled_channel = 1;
-  at::Tensor roi_actual_num = at_npu::native::OpPreparation::ApplyTensor(
-      {}, rois.options().dtype(at::kInt), rois);
-  at::Tensor x = at::ones_like(grad_input);
+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;
+  c10::SmallVector<int64_t, SIZE> xdiff_shape = at_npu::native::array_to_small_vector(grad_input.sizes());
+  int64_t roi_end_mode = 1;
+  #at::Tensor result =
+  #    OpPreparation::ApplyTensorWithFormat(grad_output, xdiff_shape, ACL_FORMAT_NC1HWC0);
   OpCommand cmd;
-  cmd.Name("RoiPoolingGradWithArgMax")
+  cmd.Name("ROIAlignGrad")
       .Input(grad_output)
-      .Input(x)
       .Input(rois)
-      .Input(roi_actual_num)
-      .Input(argmax)
-      .Output(grad_input)
-      .Attr("pooled_h", pooled_height_64)
-      .Attr("pooled_w", pooled_width_64)
-      .Attr("spatial_scale_h", spatial_scale)
-      .Attr("spatial_scale_w", spatial_scale)
-      .Attr("pool_channel", pooled_channel)
+      .Output(result)
+      .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_pool_forward_impl(Tensor input, Tensor rois, Tensor output,
-                           Tensor argmax, int pooled_height, int pooled_width,
-                           float spatial_scale);
+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_pool_backward_impl(Tensor grad_output, Tensor rois, Tensor argmax,
-                            Tensor grad_input, int pooled_height,
-                            int pooled_width, float spatial_scale);
+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_pool_forward_impl, roi_pool_forward_npu);
-REGISTER_NPU_IMPL(roi_pool_backward_impl, roi_pool_backward_npu);
+REGISTER_NPU_IMPL(roi_align_forward_impl, roi_align_forward_npu);
+REGISTER_NPU_IMPL(roi_align_backward_impl, roi_align_backward_npu);