wehub-resource-sync
244 lines
9.1 KiB
Plaintext
244 lines
9.1 KiB
Plaintext
// 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/psroi_pool_kernel.h"
|
|
|
|
#include <algorithm>
|
|
#include <vector>
|
|
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/common/place.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/tensor_utils.h"
|
|
|
|
namespace phi {
|
|
|
|
static constexpr int kNumCUDAThreads = 512;
|
|
static constexpr int kNumMaximumNumBlocks = 4096;
|
|
|
|
static inline int NumBlocks(const int64_t N) {
|
|
return std::min(static_cast<int>((N + kNumCUDAThreads - 1) / kNumCUDAThreads),
|
|
kNumMaximumNumBlocks);
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void GPUPSROIPoolForward(const int64_t nthreads,
|
|
const T* input_data,
|
|
const T* input_rois,
|
|
const float spatial_scale,
|
|
const int input_channels,
|
|
const int height,
|
|
const int width,
|
|
const int output_channels,
|
|
const int pooled_height,
|
|
const int pooled_width,
|
|
const int* rois_batch_id_data,
|
|
T* output_data) {
|
|
int64_t index =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x);
|
|
const int64_t offset = static_cast<int64_t>(blockDim.x) * gridDim.x;
|
|
for (int64_t i = index; i < nthreads; i += offset) {
|
|
// The output is in order (n, c, ph, pw)
|
|
int pw = i % pooled_width;
|
|
int ph = (i / pooled_width) % pooled_height;
|
|
int c = (i / pooled_width / pooled_height) % output_channels;
|
|
int n = i / pooled_width / pooled_height / output_channels;
|
|
|
|
// set roi_batch_id
|
|
int roi_batch_id = rois_batch_id_data[n];
|
|
|
|
// [start, end) interval for spatial sampling
|
|
const T* offset_input_rois = input_rois + n * 4;
|
|
T roi_start_w = static_cast<T>(round(offset_input_rois[0])) * spatial_scale;
|
|
T roi_start_h = static_cast<T>(round(offset_input_rois[1])) * spatial_scale;
|
|
T roi_end_w =
|
|
static_cast<T>(round(offset_input_rois[2]) + 1.) * spatial_scale;
|
|
T roi_end_h =
|
|
static_cast<T>(round(offset_input_rois[3]) + 1.) * spatial_scale;
|
|
|
|
// Force too small ROIs to be 1x1
|
|
T roi_height = max(roi_end_h - roi_start_h, (T)0.1); // avoid 0
|
|
T roi_width = max(roi_end_w - roi_start_w, (T)0.1);
|
|
|
|
// Compute w and h at input feature map
|
|
T bin_size_h = roi_height / static_cast<T>(pooled_height);
|
|
T bin_size_w = roi_width / static_cast<T>(pooled_width);
|
|
|
|
int hstart = floor(bin_size_h * static_cast<T>(ph) + roi_start_h);
|
|
int wstart = floor(bin_size_w * static_cast<T>(pw) + roi_start_w);
|
|
int hend = ceil(bin_size_h * static_cast<T>(ph + 1) + roi_start_h);
|
|
int wend = ceil(bin_size_w * static_cast<T>(pw + 1) + roi_start_w);
|
|
|
|
// Add roi offsets and clip to input boundaries
|
|
hstart = min(max(hstart, 0), height);
|
|
hend = min(max(hend, 0), height);
|
|
wstart = min(max(wstart, 0), width);
|
|
wend = min(max(wend, 0), width);
|
|
bool is_empty = (hend <= hstart) || (wend <= wstart);
|
|
|
|
int64_t input_channel =
|
|
(static_cast<int64_t>(c) * pooled_height + ph) * pooled_width + pw;
|
|
const int64_t input_offset =
|
|
(static_cast<int64_t>(roi_batch_id) * input_channels + input_channel) *
|
|
height * width;
|
|
const T* offset_input_data = input_data + input_offset;
|
|
T outsum = 0;
|
|
|
|
for (int ih = hstart; ih < hend; ++ih) {
|
|
for (int iw = wstart; iw < wend; ++iw) {
|
|
int input_index = ih * width + iw;
|
|
outsum += offset_input_data[input_index];
|
|
}
|
|
}
|
|
|
|
T bin_area = static_cast<T>((hend - hstart) * (wend - wstart));
|
|
output_data[i] = is_empty ? 0. : outsum / bin_area;
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void PsroiPoolKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& rois,
|
|
const optional<DenseTensor>& rois_num,
|
|
int pooled_height,
|
|
int pooled_width,
|
|
int output_channels,
|
|
float spatial_scale,
|
|
DenseTensor* out) {
|
|
auto in_dims = x.dims();
|
|
int batch_size = in_dims[0];
|
|
int input_channels = in_dims[1];
|
|
int height = in_dims[2];
|
|
int width = in_dims[3];
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
input_channels,
|
|
output_channels * pooled_height * pooled_width,
|
|
errors::InvalidArgument(
|
|
"The channels %d of input X should equal the product of "
|
|
"output_channels %d x pooled_height %d x pooled_width %d.",
|
|
input_channels,
|
|
output_channels,
|
|
pooled_height,
|
|
pooled_width));
|
|
|
|
int64_t rois_num_t = rois.dims()[0];
|
|
if (rois_num_t == 0) return;
|
|
int rois_batch_size;
|
|
DenseTensor rois_batch_id_list;
|
|
rois_batch_id_list.Resize({rois_num_t});
|
|
int* rois_batch_id_data =
|
|
dev_ctx.template HostAlloc<int>(&rois_batch_id_list);
|
|
|
|
if (rois_num.get_ptr()) {
|
|
rois_batch_size = rois_num->numel();
|
|
auto* rois_num_data = rois_num->data<int>();
|
|
PADDLE_ENFORCE_EQ(rois_batch_size,
|
|
batch_size,
|
|
errors::InvalidArgument(
|
|
"The batch size of input(ROIs) and input(X) must be "
|
|
"the same but received batch size of input(ROIs) and "
|
|
"input(X) is %d and %d respectively.",
|
|
rois_batch_size,
|
|
batch_size));
|
|
std::vector<int> rois_num_list(rois_batch_size);
|
|
memory_utils::Copy(CPUPlace(),
|
|
rois_num_list.data(),
|
|
dev_ctx.GetPlace(),
|
|
rois_num_data,
|
|
sizeof(int) * rois_batch_size,
|
|
0);
|
|
int rois_num_count = 0;
|
|
for (int i = 0; i < rois_batch_size; ++i) {
|
|
rois_num_count += rois_num_list[i];
|
|
}
|
|
PADDLE_ENFORCE_EQ(
|
|
rois_num_count,
|
|
rois_num_t,
|
|
errors::InvalidArgument(
|
|
"the rois_num from input and RoisNum must be the same"));
|
|
int start = 0;
|
|
for (int n = 0; n < rois_batch_size; ++n) {
|
|
for (int i = start; i < start + rois_num_list[n]; ++i) {
|
|
rois_batch_id_data[i] = n;
|
|
}
|
|
start += rois_num_list[n];
|
|
}
|
|
} else {
|
|
auto rois_lod = rois.lod().back();
|
|
rois_batch_size = rois_lod.size() - 1;
|
|
PADDLE_ENFORCE_EQ(rois_batch_size,
|
|
batch_size,
|
|
errors::InvalidArgument(
|
|
"The batch size of input(ROIs) and input(X) must be "
|
|
"the same but received batch size of input(ROIs) and "
|
|
"input(X) is %d and %d respectively.",
|
|
rois_batch_size,
|
|
batch_size));
|
|
int rois_num_with_lod = rois_lod[rois_batch_size];
|
|
PADDLE_ENFORCE_EQ(rois_num_t,
|
|
rois_num_with_lod,
|
|
errors::InvalidArgument(
|
|
"The number of rois from input(ROIs) and its LOD "
|
|
"must be the same. Received rois %d of input(ROIs), "
|
|
"but the number of rois from its LOD is %d.",
|
|
rois_num_t,
|
|
rois_num_with_lod));
|
|
|
|
// set rois batch id
|
|
for (int n = 0; n < rois_batch_size; ++n) {
|
|
for (size_t i = rois_lod[n]; i < rois_lod[n + 1]; ++i) {
|
|
rois_batch_id_data[i] = n;
|
|
}
|
|
}
|
|
}
|
|
DenseTensor rois_batch_id_list_gpu;
|
|
Copy(dev_ctx,
|
|
rois_batch_id_list,
|
|
dev_ctx.GetPlace(),
|
|
false,
|
|
&rois_batch_id_list_gpu);
|
|
|
|
int64_t output_size = out->numel();
|
|
int64_t blocks = NumBlocks(output_size);
|
|
int threads = kNumCUDAThreads;
|
|
|
|
// NOTE(large-tensor): Kernel launch requires int type for grid dimension
|
|
PADDLE_ENFORCE_LE_INT_MAX(blocks, "blocks");
|
|
// call cuda kernel function
|
|
GPUPSROIPoolForward<T>
|
|
<<<static_cast<int>(blocks), threads, 0, dev_ctx.stream()>>>(
|
|
output_size,
|
|
x.data<T>(),
|
|
rois.data<T>(),
|
|
spatial_scale,
|
|
input_channels,
|
|
height,
|
|
width,
|
|
output_channels,
|
|
pooled_height,
|
|
pooled_width,
|
|
rois_batch_id_list_gpu.data<int>(),
|
|
dev_ctx.template Alloc<T>(out));
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(
|
|
psroi_pool, GPU, ALL_LAYOUT, phi::PsroiPoolKernel, float, double) {
|
|
kernel->InputAt(2).SetDataType(phi::CppTypeToDataType<int>::Type());
|
|
}
|