Files
2026-07-13 12:40:42 +08:00

115 lines
3.9 KiB
C++

// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "paddle/phi/kernels/roi_pool_grad_kernel.h"
#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/empty_kernel.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/funcs/math_function.h"
namespace phi {
template <typename T, typename Context>
void RoiPoolGradKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& boxes,
const optional<DenseTensor>& boxes_num,
const DenseTensor& arg_max,
const DenseTensor& out_grad,
int pooled_height,
int pooled_width,
float spatial_scale,
DenseTensor* dx) {
if (x.numel() == 0 || boxes.numel() == 0) {
Full<T, Context>(dev_ctx, dx->dims(), 0, dx);
return;
}
if (dx) {
int rois_num = static_cast<int>(boxes.dims()[0]);
DenseTensor box_batch_id_list = Empty<int>(dev_ctx, {rois_num});
int* box_batch_id_data = box_batch_id_list.data<int>();
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;
}
start += boxes_num_data[n];
}
} else {
auto boxes_lod = boxes.lod().back();
boxes_batch_size = static_cast<int>(boxes_lod.size() - 1);
for (int n = 0; n < boxes_batch_size; ++n) {
for (size_t i = boxes_lod[n]; i < boxes_lod[n + 1]; ++i) {
box_batch_id_data[i] = n;
}
}
}
const T* boxes_data = boxes.data<T>();
const T* out_grad_data = out_grad.data<T>();
const int64_t* arg_max_data = arg_max.data<int64_t>();
T* dx_data = dev_ctx.template Alloc<T>(dx);
funcs::SetConstant<Context, T> set_zero;
set_zero(dev_ctx, dx, static_cast<T>(0));
auto in_stride = common::stride(x.dims());
auto arg_max_stride = common::stride(arg_max.dims());
auto roi_stride = common::stride(boxes.dims());
auto out_stride = common::stride(out_grad.dims());
int channels = static_cast<int>(x.dims()[1]);
for (int n = 0; n < rois_num; ++n) {
int roi_batch_idx = box_batch_id_data[n];
T* batch_grad_data = dx_data + roi_batch_idx * in_stride[0];
for (int c = 0; c < channels; ++c) {
for (int ph = 0; ph < pooled_height; ++ph) {
for (int pw = 0; pw < pooled_width; ++pw) {
int pool_index = ph * pooled_width + pw;
if (arg_max_data[pool_index] >= 0) {
auto index = arg_max_data[pool_index];
batch_grad_data[index] += out_grad_data[pool_index];
}
}
}
batch_grad_data += in_stride[1];
out_grad_data += out_stride[1];
arg_max_data += arg_max_stride[1];
}
boxes_data += roi_stride[0];
}
}
}
} // namespace phi
PD_REGISTER_KERNEL(roi_pool_grad,
CPU,
ALL_LAYOUT,
phi::RoiPoolGradKernel,
float,
double,
int) {
kernel->InputAt(3).SetDataType(phi::DataType::INT64);
}