roi_align_grad boxes_num support int64 (#70003)

wanghuancoder · web-flow · commit 7a9149517da9 · 2024-12-06T12:53:43.000+08:00
diff --git a/paddle/phi/kernels/cpu/roi_align_grad_kernel.cc b/paddle/phi/kernels/cpu/roi_align_grad_kernel.cc
@@ -97,13 +97,24 @@ void RoiAlignGradKernel(const Context& dev_ctx,
   int boxes_batch_size = 0;
   if (boxes_num) {
     boxes_batch_size = static_cast<int>(boxes_num->numel());
-    auto* boxes_num_data = boxes_num->data<int>();
-    int start = 0;
-    for (int n = 0; n < boxes_batch_size; ++n) {
-      for (int i = start; i < start + boxes_num_data[n]; ++i) {
-        box_batch_id_data[i] = n;
+    if (boxes_num->dtype() == phi::DataType::INT64) {
+      auto* boxes_num_data = boxes_num->data<int64_t>();
+      int64_t start = 0;
+      for (int64_t n = 0; n < boxes_batch_size; ++n) {
+        for (int64_t i = start; i < start + boxes_num_data[n]; ++i) {
+          box_batch_id_data[i] = n;
+        }
+        start += boxes_num_data[n];
+      }
+    } else if (boxes_num->dtype() == phi::DataType::INT32) {
+      auto* boxes_num_data = boxes_num->data<int>();
+      int start = 0;
+      for (int n = 0; n < boxes_batch_size; ++n) {
+        for (int i = start; i < start + boxes_num_data[n]; ++i) {
+          box_batch_id_data[i] = n;
+        }
+        start += boxes_num_data[n];
       }
-      start += boxes_num_data[n];
     }
   } else {
     auto boxes_lod = boxes.lod().back();
diff --git a/paddle/phi/kernels/gpu/roi_align_grad_kernel.cu b/paddle/phi/kernels/gpu/roi_align_grad_kernel.cu
@@ -193,19 +193,36 @@ void RoiAlignGradKernel(const Context& dev_ctx,
   auto gplace = dev_ctx.GetPlace();
   if (boxes_num) {
     int boxes_batch_size = boxes_num->numel();
-    std::vector<int> boxes_num_list(boxes_batch_size);
-    memory_utils::Copy(cplace,
-                       boxes_num_list.data(),
-                       gplace,
-                       boxes_num->data<int>(),
-                       sizeof(int) * boxes_batch_size,
-                       0);
-    int start = 0;
-    for (int n = 0; n < boxes_batch_size; ++n) {
-      for (size_t i = start; i < start + boxes_num_list[n]; ++i) {
-        box_batch_size[i] = n;
+    if (boxes_num->dtype() == phi::DataType::INT64) {
+      std::vector<int64_t> boxes_num_list(boxes_batch_size);
+      memory_utils::Copy(cplace,
+                         boxes_num_list.data(),
+                         gplace,
+                         boxes_num->data<int64_t>(),
+                         sizeof(int64_t) * boxes_batch_size,
+                         0);
+      int64_t start = 0;
+      for (int64_t n = 0; n < boxes_batch_size; ++n) {
+        for (int64_t i = start; i < start + boxes_num_list[n]; ++i) {
+          box_batch_size[i] = n;
+        }
+        start += boxes_num_list[n];
+      }
+    } else if (boxes_num->dtype() == phi::DataType::INT32) {
+      std::vector<int> boxes_num_list(boxes_batch_size);
+      memory_utils::Copy(cplace,
+                         boxes_num_list.data(),
+                         gplace,
+                         boxes_num->data<int>(),
+                         sizeof(int) * boxes_batch_size,
+                         0);
+      int start = 0;
+      for (int n = 0; n < boxes_batch_size; ++n) {
+        for (size_t i = start; i < start + boxes_num_list[n]; ++i) {
+          box_batch_size[i] = n;
+        }
+        start += boxes_num_list[n];
       }
-      start += boxes_num_list[n];
     }
   } else {
     auto boxes_lod = boxes.lod().back();
diff --git a/paddle/phi/kernels/xpu/roi_align_grad_kernel.cc b/paddle/phi/kernels/xpu/roi_align_grad_kernel.cc
@@ -50,16 +50,30 @@ void RoiAlignGradKernel(const Context& dev_ctx,
   int* cpu_lod = nullptr;
   if (boxes_num) {
     rois_batch_size = boxes_num->numel();
-    std::vector<int> rois_num_list(rois_batch_size);
-    memory_utils::Copy(cplace,
-                       rois_num_list.data(),
-                       xplace,
-                       boxes_num->data<int>(),
-                       sizeof(int) * rois_batch_size);
-    cpu_lod = new int[rois_batch_size + 1];
-    cpu_lod[0] = 0;
-    for (int i = 0; i < rois_batch_size; i++) {
-      cpu_lod[i + 1] = cpu_lod[i] + rois_num_list[i];
+    if (boxes_num->dtype() == phi::DataType::INT64) {
+      std::vector<int64_t> rois_num_list(rois_batch_size);
+      memory_utils::Copy(cplace,
+                         rois_num_list.data(),
+                         xplace,
+                         boxes_num->data<int64_t>(),
+                         sizeof(int64_t) * rois_batch_size);
+      cpu_lod = new int[rois_batch_size + 1];
+      cpu_lod[0] = 0;
+      for (int64_t i = 0; i < rois_batch_size; i++) {
+        cpu_lod[i + 1] = cpu_lod[i] + rois_num_list[i];
+      }
+    } else if (boxes_num->dtype() == phi::DataType::INT32) {
+      std::vector<int> rois_num_list(rois_batch_size);
+      memory_utils::Copy(cplace,
+                         rois_num_list.data(),
+                         xplace,
+                         boxes_num->data<int>(),
+                         sizeof(int) * rois_batch_size);
+      cpu_lod = new int[rois_batch_size + 1];
+      cpu_lod[0] = 0;
+      for (int i = 0; i < rois_batch_size; i++) {
+        cpu_lod[i + 1] = cpu_lod[i] + rois_num_list[i];
+      }
     }
   } else {
     auto rois_lod = boxes.lod().back();
