wehub-resource-sync
694 lines
26 KiB
Plaintext
694 lines
26 KiB
Plaintext
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*
|
|
* 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 "NvInfer.h"
|
|
#include "common/plugin.h"
|
|
#include <algorithm>
|
|
#include <assert.h>
|
|
#include <cub/cub.cuh>
|
|
#include <cuda.h>
|
|
#include <cuda_runtime.h>
|
|
#include <functional>
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <vector>
|
|
|
|
namespace nvinfer1
|
|
{
|
|
namespace plugin
|
|
{
|
|
#define PLUGIN_CHECK_CUDA(call) \
|
|
do \
|
|
{ \
|
|
cudaError_t status = call; \
|
|
if (status != cudaSuccess) \
|
|
{ \
|
|
return status; \
|
|
} \
|
|
} while (0)
|
|
template <typename TFloat>
|
|
struct Bbox
|
|
{
|
|
TFloat x1, y1, x2, y2;
|
|
};
|
|
|
|
typedef nvinfer1::DataType DType_t;
|
|
|
|
typedef enum
|
|
{
|
|
NCHW = 0,
|
|
NC4HW = 1
|
|
} DLayout_t;
|
|
|
|
typedef pluginStatus_t frcnnStatus_t;
|
|
|
|
#define DEBUG_RPN_ENABLE 0
|
|
|
|
#define FRCNN_ASSERT_PARAM(exp) \
|
|
do \
|
|
{ \
|
|
if (!(exp)) \
|
|
{ \
|
|
DEBUG_FPRINTF(stderr, "Bad param - " #exp ", %s:%d\n", __FILE__, __LINE__); \
|
|
return STATUS_BAD_PARAM; \
|
|
} \
|
|
} while (0)
|
|
|
|
#define DEBUG_FPRINTF(...) \
|
|
do \
|
|
{ \
|
|
if (DEBUG_RPN_ENABLE) \
|
|
{ \
|
|
fprintf(__VA_ARGS__); \
|
|
} \
|
|
} while (0)
|
|
|
|
#define CUDA_MEM_ALIGN 256
|
|
|
|
unsigned int hash(const void* array_, size_t size);
|
|
int8_t* alignPtr(int8_t* ptr, uintptr_t to);
|
|
__global__ void setOffset(int stride, int size, int* output);
|
|
frcnnStatus_t nms(cudaStream_t stream,
|
|
const int N,
|
|
const int R,
|
|
const int preNmsTop,
|
|
const int nmsMaxOut,
|
|
const float iouThreshold,
|
|
const DType_t t_fgScores,
|
|
const DLayout_t l_fgScores,
|
|
void* fgScores,
|
|
const DType_t t_proposals,
|
|
const DLayout_t l_proposals,
|
|
const void* proposals,
|
|
void* workspace,
|
|
const DType_t t_rois,
|
|
void* rois);
|
|
int8_t* nextWorkspacePtr(int8_t* ptr, uintptr_t previousWorkspaceSize);
|
|
|
|
|
|
template <typename TFloat>
|
|
__device__ __host__ inline float IoU(const Bbox<TFloat>& a, const Bbox<TFloat>& b)
|
|
{
|
|
TFloat left = max(a.x1, b.x1), right = min(a.x2, b.x2);
|
|
TFloat top = max(a.y1, b.y1), bottom = min(a.y2, b.y2);
|
|
TFloat width = max((TFloat)(right - left + (TFloat) 1.0), (TFloat) 0.0);
|
|
TFloat height = max((TFloat)(bottom - top + (TFloat) 1.0), (TFloat) 0.0);
|
|
TFloat interS = width * height;
|
|
TFloat Sa = (a.x2 - a.x1 + (TFloat) 1) * (a.y2 - a.y1 + (TFloat) 1);
|
|
TFloat Sb = (b.x2 - b.x1 + (TFloat) 1) * (b.y2 - b.y1 + (TFloat) 1);
|
|
return (float) interS / (float) (Sa + Sb - interS);
|
|
}
|
|
|
|
|
|
// NMS KERNEL FOR SMALL BATCH SIZE {{{
|
|
template <typename T_PROPOSALS, typename T_ROIS, int DIM, int TSIZE>
|
|
__global__ __launch_bounds__(DIM) void nmsKernel1(const int propSize,
|
|
Bbox<T_PROPOSALS> const* __restrict__ preNmsProposals,
|
|
T_ROIS* __restrict__ afterNmsProposals,
|
|
const int preNmsTopN,
|
|
const float nmsThres,
|
|
const int afterNmsTopN)
|
|
{
|
|
__shared__ bool kept_boxes[TSIZE * DIM];
|
|
int kept = 0;
|
|
int batch_offset = blockIdx.x * propSize;
|
|
int max_box_idx = batch_offset + preNmsTopN;
|
|
int batch_offset_out = blockIdx.x * afterNmsTopN;
|
|
int flag_idx[TSIZE];
|
|
int boxes_idx[TSIZE];
|
|
Bbox<T_PROPOSALS> cur_boxes[TSIZE];
|
|
// initialize kept_boxes
|
|
#pragma unroll
|
|
|
|
for (int i = 0; i < TSIZE; i++)
|
|
{
|
|
boxes_idx[i] = threadIdx.x + batch_offset + DIM * i;
|
|
flag_idx[i] = threadIdx.x + DIM * i;
|
|
|
|
if (boxes_idx[i] < max_box_idx)
|
|
{
|
|
cur_boxes[i] = preNmsProposals[boxes_idx[i]];
|
|
kept_boxes[flag_idx[i]] = true;
|
|
}
|
|
else
|
|
{
|
|
kept_boxes[flag_idx[i]] = false;
|
|
boxes_idx[i] = -1.0f;
|
|
flag_idx[i] = -1.0f;
|
|
}
|
|
}
|
|
|
|
int ref_box_idx = 0 + batch_offset;
|
|
|
|
// remove the overlapped boxes
|
|
while ((kept < afterNmsTopN) && (ref_box_idx < max_box_idx))
|
|
{
|
|
Bbox<T_PROPOSALS> ref_box;
|
|
ref_box = preNmsProposals[ref_box_idx];
|
|
#pragma unroll
|
|
|
|
for (int i = 0; i < TSIZE; i++)
|
|
{
|
|
if (boxes_idx[i] > ref_box_idx)
|
|
{
|
|
if (IoU(ref_box, cur_boxes[i]) > nmsThres)
|
|
{
|
|
kept_boxes[flag_idx[i]] = false;
|
|
}
|
|
}
|
|
else if (boxes_idx[i] == ref_box_idx)
|
|
{
|
|
afterNmsProposals[(batch_offset_out + kept) * 4 + 0] = ref_box.x1;
|
|
afterNmsProposals[(batch_offset_out + kept) * 4 + 1] = ref_box.y1;
|
|
afterNmsProposals[(batch_offset_out + kept) * 4 + 2] = ref_box.x2;
|
|
afterNmsProposals[(batch_offset_out + kept) * 4 + 3] = ref_box.y2;
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
do
|
|
{
|
|
ref_box_idx++;
|
|
}
|
|
while (!kept_boxes[ref_box_idx - batch_offset] && ref_box_idx < max_box_idx);
|
|
|
|
kept++;
|
|
}
|
|
}
|
|
// }}}
|
|
|
|
// NMS KERNEL FOR LARGE BATCH SIZE {{{
|
|
template <typename T_PROPOSALS, typename T_ROIS, int DIM, int TSIZE>
|
|
__global__ __launch_bounds__(DIM) void nmsKernel2(const int propSize,
|
|
Bbox<T_PROPOSALS> const* __restrict__ proposals,
|
|
T_ROIS* __restrict__ filtered,
|
|
const int preNmsTopN,
|
|
const float nmsThres,
|
|
const int afterNmsTopN)
|
|
{
|
|
Bbox<T_PROPOSALS> const* cProposals = proposals + blockIdx.x * propSize;
|
|
Bbox<T_PROPOSALS> t[TSIZE];
|
|
uint64_t del = 0;
|
|
|
|
for (int i = 0; i < TSIZE; i++)
|
|
{
|
|
if (i < TSIZE - 1 || i * DIM + threadIdx.x < preNmsTopN)
|
|
{
|
|
t[i] = cProposals[i * DIM + threadIdx.x];
|
|
}
|
|
}
|
|
|
|
__shared__ Bbox<T_PROPOSALS> last;
|
|
__shared__ bool kept;
|
|
__shared__ int foundBatch;
|
|
|
|
if (threadIdx.x == 0)
|
|
{
|
|
foundBatch = 0;
|
|
}
|
|
|
|
for (int i = 0; i < TSIZE; i++)
|
|
{
|
|
for (int j = 0; j < DIM; j++)
|
|
{
|
|
int offset = i * DIM;
|
|
int index = offset + j;
|
|
|
|
if (index >= preNmsTopN)
|
|
{
|
|
break;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x == j)
|
|
{
|
|
kept = 0 == (del & ((uint64_t) 1 << i));
|
|
last = t[i];
|
|
|
|
if (kept)
|
|
{
|
|
int cnt = blockIdx.x * afterNmsTopN + foundBatch;
|
|
filtered[cnt * 4 + 0] = t[i].x1;
|
|
filtered[cnt * 4 + 1] = t[i].y1;
|
|
filtered[cnt * 4 + 2] = t[i].x2;
|
|
filtered[cnt * 4 + 3] = t[i].y2;
|
|
foundBatch++;
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if (foundBatch == afterNmsTopN)
|
|
{
|
|
return;
|
|
}
|
|
|
|
if (kept)
|
|
{
|
|
Bbox<T_PROPOSALS> test = last;
|
|
|
|
for (int k = 0; k < TSIZE; k++)
|
|
{
|
|
if (index < k * DIM + threadIdx.x
|
|
&& IoU<T_PROPOSALS>(test, t[k]) > nmsThres)
|
|
{
|
|
del |= (uint64_t) 1 << k;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// }}}
|
|
|
|
// NMS LAUNCH {{{
|
|
template <typename T_PROPOSALS, DLayout_t L_PROPOSALS, typename T_ROIS>
|
|
frcnnStatus_t nmsLaunch(cudaStream_t stream,
|
|
const int batch,
|
|
const int propSize,
|
|
void* proposals,
|
|
void* filtered,
|
|
const int preNmsTopN,
|
|
const float nmsThres,
|
|
const int afterNmsTopN)
|
|
{
|
|
const int blockSize = 1024;
|
|
#define P1(tsize) nmsKernel1<T_PROPOSALS, T_ROIS, blockSize, (tsize)>
|
|
#define P2(tsize) nmsKernel2<T_PROPOSALS, T_ROIS, blockSize, (tsize)>
|
|
void (*kernel[64])(int, Bbox<T_PROPOSALS> const*, T_ROIS*, int, float, int) =
|
|
{
|
|
P1(1), P1(2), P1(3), P1(4), P1(5), P1(6), P1(7), P1(8), P1(9), P1(10), P1(11), P1(12), P2(13), P2(14), P2(15), P2(16),
|
|
P2(17), P2(18), P2(19), P2(20), P2(21), P2(22), P2(23), P2(24), P2(25), P2(26), P2(27), P2(28), P2(29), P2(30), P2(31), P2(32),
|
|
P2(33), P2(34), P2(35), P2(36), P2(37), P2(38), P2(39), P2(40), P2(41), P2(42), P2(43), P2(44), P2(45), P2(46), P2(47), P2(48),
|
|
P2(49), P2(50), P2(51), P2(52), P2(53), P2(54), P2(55), P2(56), P2(57), P2(58), P2(59), P2(60), P2(61), P2(62), P2(63), P2(64)
|
|
};
|
|
FRCNN_ASSERT_PARAM(preNmsTopN < 64 * blockSize);
|
|
CSC(cudaMemsetAsync(filtered, 0, batch * afterNmsTopN * 4 * sizeof(T_ROIS), stream),
|
|
STATUS_FAILURE);
|
|
kernel[(preNmsTopN + blockSize - 1) / blockSize - 1] <<< batch, blockSize, 0, stream>>>(propSize,
|
|
(Bbox<T_PROPOSALS>*) proposals,
|
|
(T_ROIS*) filtered,
|
|
preNmsTopN,
|
|
nmsThres,
|
|
afterNmsTopN);
|
|
CSC(cudaGetLastError(), STATUS_FAILURE);
|
|
return STATUS_SUCCESS;
|
|
}
|
|
// }}}
|
|
|
|
// NMS GPU {{{
|
|
template <typename T_SCORES, typename T_ROIS>
|
|
frcnnStatus_t nmsGpu(cudaStream_t stream, const int N, const int R, const int preNmsTop, const int nmsMaxOut,
|
|
const float iouThreshold, void* fgScores, const void* proposals, void* workspace, void* rois)
|
|
{
|
|
// CUB's bug workaround:
|
|
// To work properly for large batch size CUB segmented sort needs ridiculous
|
|
// workspace alignment.
|
|
constexpr uintptr_t kALIGNMENT = 1 << 20;
|
|
|
|
int8_t* vworkspace = alignPtr((int8_t*) workspace, kALIGNMENT);
|
|
DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposals, N * R * 4 * sizeof(float)));
|
|
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(fgScores, N * R * sizeof(float)));
|
|
frcnnStatus_t error;
|
|
DEBUG_PRINTF("&&&& [NMS] DISCARD\n");
|
|
DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposals, N * R * 4 * sizeof(float)));
|
|
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(fgScores, N * R * sizeof(float)));
|
|
// Generate offsets
|
|
int* offsets = (int*) vworkspace;
|
|
setOffset<<<1, 1024, 0, stream>>>(R, N + 1, offsets);
|
|
CSC(cudaGetLastError(), STATUS_FAILURE);
|
|
vworkspace = vworkspace + N + 1;
|
|
vworkspace = alignPtr(vworkspace, kALIGNMENT);
|
|
// Sort (batched)
|
|
std::size_t tempStorageBytes = 0;
|
|
cub::DeviceSegmentedRadixSort::SortPairsDescending(NULL, tempStorageBytes, (T_SCORES*) fgScores,
|
|
(T_SCORES*) fgScores, (Bbox<T_ROIS>*) proposals, (Bbox<T_ROIS>*) proposals, N * R, N, offsets, offsets + 1, 0,
|
|
8 * sizeof(T_SCORES), stream);
|
|
CSC(cudaGetLastError(), STATUS_FAILURE);
|
|
T_SCORES* scoresOut = (T_SCORES*) vworkspace;
|
|
vworkspace = (int8_t*) (scoresOut + N * R);
|
|
vworkspace = alignPtr(vworkspace, kALIGNMENT);
|
|
Bbox<T_ROIS>* proposalsOut = (Bbox<T_ROIS>*) vworkspace;
|
|
vworkspace = (int8_t*) (proposalsOut + N * R);
|
|
vworkspace = alignPtr(vworkspace, kALIGNMENT);
|
|
cub::DeviceSegmentedRadixSort::SortPairsDescending(vworkspace, tempStorageBytes, (T_SCORES*) fgScores,
|
|
(T_SCORES*) scoresOut, (Bbox<T_ROIS>*) proposals, (Bbox<T_ROIS>*) proposalsOut, N * R, N, offsets, offsets + 1,
|
|
0, 8 * sizeof(T_SCORES), stream);
|
|
CSC(cudaGetLastError(), STATUS_FAILURE);
|
|
DEBUG_PRINTF("&&&& [NMS] POST CUB\n");
|
|
DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposalsOut, N * R * 4 * sizeof(float)));
|
|
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(scoresOut, N * R * sizeof(float)));
|
|
error = nmsLaunch<T_ROIS, NC4HW, T_ROIS>(stream,
|
|
N,
|
|
R,
|
|
proposalsOut,
|
|
rois,
|
|
preNmsTop,
|
|
iouThreshold,
|
|
nmsMaxOut);
|
|
DEBUG_PRINTF("&&&& [NMS] POST LAUNCH\n");
|
|
DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(rois, N * nmsMaxOut * 4 * sizeof(float)));
|
|
|
|
if (error != STATUS_SUCCESS)
|
|
{
|
|
return error;
|
|
}
|
|
|
|
return STATUS_SUCCESS;
|
|
}
|
|
// }}}
|
|
|
|
typedef frcnnStatus_t (*nmsFun)(cudaStream_t,
|
|
const int, // N
|
|
const int, // R
|
|
const int, // preNmsTop
|
|
const int, // nmsMaxOut
|
|
const float, // iouThreshold
|
|
void*, // fgScores
|
|
const void*, // proposals,
|
|
void*, // workspace,
|
|
void*); // rois
|
|
|
|
struct nmsLaunchConfig
|
|
{
|
|
DType_t t_fgScores;
|
|
DLayout_t l_fgScores;
|
|
DType_t t_proposals;
|
|
DLayout_t l_proposals;
|
|
DType_t t_rois;
|
|
nmsFun function;
|
|
|
|
nmsLaunchConfig(DType_t t_fgScores,
|
|
DLayout_t l_fgScores,
|
|
DType_t t_proposals,
|
|
DLayout_t l_proposals,
|
|
DType_t t_rois,
|
|
nmsFun function)
|
|
: t_fgScores(t_fgScores)
|
|
, l_fgScores(l_fgScores)
|
|
, t_proposals(t_proposals)
|
|
, l_proposals(l_proposals)
|
|
, t_rois(t_rois)
|
|
, function(function)
|
|
{
|
|
}
|
|
|
|
nmsLaunchConfig(
|
|
DType_t t_fgScores, DLayout_t l_fgScores, DType_t t_proposals, DLayout_t l_proposals, DType_t t_rois)
|
|
: t_fgScores(t_fgScores)
|
|
, l_fgScores(l_fgScores)
|
|
, t_proposals(t_proposals)
|
|
, l_proposals(l_proposals)
|
|
, t_rois(t_rois)
|
|
, function(nullptr)
|
|
{
|
|
}
|
|
|
|
bool operator==(nmsLaunchConfig const& other) const
|
|
{
|
|
return (t_fgScores == other.t_fgScores) && (l_fgScores == other.l_fgScores)
|
|
&& (t_proposals == other.t_proposals) && (l_proposals == other.l_proposals)
|
|
&& (t_rois == other.t_rois);
|
|
}
|
|
};
|
|
|
|
static std::vector<nmsLaunchConfig> nmsLCVec;
|
|
#define FLOAT32 nvinfer1::DataType::kFLOAT
|
|
|
|
__global__ void _inverse_transform_gpu(const float* RPN_prob, const float* RPN_regr, int N,
|
|
int INPUT_H, int INPUT_W, int RPN_H, int RPN_W, float RPN_STD_SCALING, int RPN_STRIDE,
|
|
float* ANCHOR_SIZES, int anc_size_num, float* ANCHOR_RATIOS, int anc_ratio_num, float bbox_min_size,
|
|
float* fg_scores, float* proposal_out)
|
|
{
|
|
int nthreads = N * RPN_H * RPN_W * anc_size_num * anc_ratio_num;
|
|
int num_ancs = anc_size_num * anc_ratio_num;
|
|
|
|
for (int out_idx = threadIdx.x + blockDim.x * blockIdx.x; out_idx < nthreads;
|
|
out_idx += blockDim.x * gridDim.x)
|
|
{
|
|
//input RPN_regr: (N, A4, H, W), thread: (N, A, H, W)
|
|
int idx = out_idx;
|
|
int w = idx % RPN_W;
|
|
idx /= RPN_W;
|
|
int h = idx % RPN_H;
|
|
idx /= RPN_H;
|
|
int a = idx % num_ancs;
|
|
int n = idx / num_ancs;
|
|
// normalize by RPN_STD_SCALING
|
|
int ptr_1 = ((((n * num_ancs) + a) * 4) * RPN_H + h) * RPN_W + w;
|
|
int ptr_2 = ((((n * num_ancs) + a) * 4 + 1) * RPN_H + h) * RPN_W + w;
|
|
int ptr_3 = ((((n * num_ancs) + a) * 4 + 2) * RPN_H + h) * RPN_W + w;
|
|
int ptr_4 = ((((n * num_ancs) + a) * 4 + 3) * RPN_H + h) * RPN_W + w;
|
|
float tx = RPN_regr[ptr_1] / RPN_STD_SCALING;
|
|
float ty = RPN_regr[ptr_2] / RPN_STD_SCALING;
|
|
float tw = RPN_regr[ptr_3] / RPN_STD_SCALING;
|
|
float th = RPN_regr[ptr_4] / RPN_STD_SCALING;
|
|
// do inverse transform
|
|
int ar = a % anc_ratio_num;
|
|
int as = a / anc_ratio_num;
|
|
float anchor_w = ANCHOR_SIZES[as] * ANCHOR_RATIOS[ar];
|
|
float anchor_h = ANCHOR_SIZES[as] / ANCHOR_RATIOS[ar];
|
|
float anchor_cx = (w + 0.5f) * RPN_STRIDE;
|
|
float anchor_cy = (h + 0.5f) * RPN_STRIDE;
|
|
float cx1 = anchor_cx + anchor_w * tx;
|
|
float cy1 = anchor_cy + anchor_h * ty;
|
|
float w1 = __expf(tw) * anchor_w;
|
|
float h1 = __expf(th) * anchor_h;
|
|
tx = cx1 - w1 / 2.0f;
|
|
ty = cy1 - h1 / 2.0f;
|
|
tw = w1;
|
|
th = h1;
|
|
tw += tx;
|
|
th += ty;
|
|
// clip to min
|
|
tx = (tx >= 0.0f) ? tx : 0.0f;
|
|
ty = (ty >= 0.0f) ? ty : 0.0f;
|
|
tw = (tw >= 0.0f) ? tw : 0.0f;
|
|
th = (th >= 0.0f) ? th : 0.0f;
|
|
//clip to max
|
|
tx = (tx <= INPUT_W - 1.0f) ? tx : (INPUT_W - 1.0f);
|
|
ty = (ty <= INPUT_H - 1.0f) ? ty : (INPUT_H - 1.0f);
|
|
tw = (tw <= INPUT_W - 1.0f) ? tw : (INPUT_W - 1.0f);
|
|
th = (th <= INPUT_H - 1.0f) ? th : (INPUT_H - 1.0f);
|
|
// filter out small boxes by setting the confidence to -inf
|
|
int ininf = 0xff800000;
|
|
float ninf = *(float*) &ininf;
|
|
|
|
if (tw - tx <= bbox_min_size || th - ty <= bbox_min_size)
|
|
{
|
|
fg_scores[out_idx] = ninf;
|
|
}
|
|
|
|
// copy to proposal_out, output shape: (N, A, H, W, 4)
|
|
proposal_out[out_idx * 4] = tx;
|
|
proposal_out[out_idx * 4 + 1] = ty;
|
|
proposal_out[out_idx * 4 + 2] = tw;
|
|
proposal_out[out_idx * 4 + 3] = th;
|
|
}
|
|
}
|
|
|
|
|
|
|
|
cudaError_t _inverse_transform_wrapper(const float* RPN_prob, const float* RPN_regr, int N, int INPUT_H,
|
|
int INPUT_W, int RPN_H, int RPN_W, float RPN_STD_SCALING, int RPN_STRIDE, float* ANCHOR_SIZES,
|
|
int anc_size_num, float* ANCHOR_RATIOS, int anc_ratio_num, float bbox_min_size, float* fg_scores,
|
|
float* proposal_out, cudaStream_t stream)
|
|
{
|
|
const int block_size = 1024;
|
|
const int grid_size = (N * anc_size_num * anc_ratio_num * RPN_H * RPN_W + block_size - 1) /
|
|
(block_size);
|
|
_inverse_transform_gpu <<< grid_size, block_size, 0, stream>>> (RPN_prob, RPN_regr, N, INPUT_H,
|
|
INPUT_W, RPN_H, RPN_W, RPN_STD_SCALING, RPN_STRIDE, ANCHOR_SIZES, anc_size_num, ANCHOR_RATIOS,
|
|
anc_ratio_num, bbox_min_size, fg_scores, proposal_out);
|
|
|
|
return cudaGetLastError();
|
|
}
|
|
|
|
size_t _proposalsForwardNMSWorkspaceSize(int N,
|
|
int A,
|
|
int H,
|
|
int W,
|
|
int nmsMaxOut)
|
|
{
|
|
return N * A * H * W * 5 * 5 * sizeof(float) + (1 << 22);
|
|
}
|
|
|
|
size_t _proposalsForwardBboxWorkspaceSize(int N, int A, int H, int W)
|
|
{
|
|
return N * A * H * W * 4 * sizeof(float);
|
|
}
|
|
|
|
|
|
size_t _proposalForwardFgScoresWorkspaceSize(int N, int A, int H, int W)
|
|
{
|
|
return N * A * H * W * sizeof(float);
|
|
}
|
|
|
|
|
|
size_t anchors_buf_size(int anc_size_num, int anc_ratio_num)
|
|
{
|
|
return (anc_size_num + anc_ratio_num) * sizeof(float);
|
|
}
|
|
|
|
size_t calculateTotalWorkspaceSize(size_t* workspaces, int count);
|
|
|
|
size_t _get_workspace_size(int N,
|
|
int anc_size_num,
|
|
int anc_ratio_num,
|
|
int H,
|
|
int W,
|
|
int nmsMaxOut)
|
|
{
|
|
size_t wss[4];
|
|
int A = anc_size_num * anc_ratio_num;
|
|
wss[0] = _proposalsForwardNMSWorkspaceSize(N, A, H, W, nmsMaxOut);
|
|
wss[1] = _proposalsForwardBboxWorkspaceSize(N, A, H, W);
|
|
wss[2] = _proposalForwardFgScoresWorkspaceSize(N, A, H, W);
|
|
wss[3] = anchors_buf_size(anc_size_num, anc_ratio_num);
|
|
return calculateTotalWorkspaceSize(wss, 4);
|
|
}
|
|
|
|
|
|
|
|
template <typename T>
|
|
frcnnStatus_t extractFgScores_gpu(cudaStream_t stream,
|
|
int N,
|
|
int A,
|
|
int H,
|
|
int W,
|
|
const void* scores,
|
|
void* fgScores)
|
|
{
|
|
//TODO custom kernel for this
|
|
size_t size = A * H * W * sizeof(T);
|
|
|
|
for (int n = 0; n < N; n++)
|
|
{
|
|
size_t offset_ld = n * A * H * W;
|
|
size_t offset_st = n * A * H * W;
|
|
CSC(cudaMemcpyAsync(((T*) fgScores) + offset_st, ((T*) scores) + offset_ld, size,
|
|
cudaMemcpyDeviceToDevice, stream), STATUS_FAILURE);
|
|
}
|
|
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
cudaError_t _copy_anchors_to_gpu(cudaStream_t stream, float* ANCHOR_SIZES, int anc_size_num, float* ANCHOR_RATIOS,
|
|
int anc_ratio_num, void* anchor_size_buf)
|
|
{
|
|
PLUGIN_CHECK_CUDA(cudaMemcpyAsync(anchor_size_buf, static_cast<void*>(ANCHOR_SIZES), sizeof(float) * anc_size_num,
|
|
cudaMemcpyHostToDevice, stream));
|
|
PLUGIN_CHECK_CUDA(cudaMemcpyAsync(static_cast<void*>(static_cast<float*>(anchor_size_buf) + anc_size_num),
|
|
static_cast<void*>(ANCHOR_RATIOS), sizeof(float) * anc_ratio_num, cudaMemcpyHostToDevice, stream));
|
|
|
|
return cudaSuccess;
|
|
}
|
|
|
|
|
|
__global__ void _normalize_rois_kernel(float* roi_after_nms, int nthreads, int width, int height)
|
|
{
|
|
for(int i = threadIdx.x + blockDim.x * blockIdx.x; i < nthreads; i += blockDim.x * gridDim.x)
|
|
{
|
|
float x1 = roi_after_nms[i * 4];
|
|
float y1 = roi_after_nms[i * 4 + 1];
|
|
float x2 = roi_after_nms[i * 4 + 2];
|
|
float y2 = roi_after_nms[i * 4 + 3];
|
|
roi_after_nms[i * 4] = y1 / (height - 1.0f);
|
|
roi_after_nms[i * 4 + 1] = x1 / (width - 1.0f);
|
|
roi_after_nms[i * 4 + 2] = y2 / (height - 1.0f);
|
|
roi_after_nms[i * 4 + 3] = x2 / (width - 1.0f);
|
|
}
|
|
}
|
|
|
|
|
|
|
|
cudaError_t _normalize_rois(float* roi_after_nms, int n, int max_box_num, int input_width,
|
|
int input_height, cudaStream_t stream)
|
|
{
|
|
const int block_size = 1024;
|
|
const int grid_size = (n * max_box_num + block_size - 1) / block_size;
|
|
_normalize_rois_kernel <<< grid_size, block_size, 0, stream>>>(roi_after_nms, n * max_box_num,
|
|
input_width, input_height);
|
|
|
|
return cudaGetLastError();
|
|
}
|
|
|
|
|
|
int proposalInference_gpu(
|
|
cudaStream_t stream,
|
|
const void* rpn_prob,
|
|
const void* rpn_regr,
|
|
int batch_size,
|
|
int input_height,
|
|
int input_width,
|
|
int rpn_height,
|
|
int rpn_width,
|
|
int MAX_BOX_NUM,
|
|
int RPN_PRE_NMS_TOP_N,
|
|
float* ANCHOR_SIZES,
|
|
int anc_size_num,
|
|
float* ANCHOR_RATIOS,
|
|
int anc_ratio_num,
|
|
float rpn_std_scaling,
|
|
int rpn_stride,
|
|
float bbox_min_size,
|
|
float nms_iou_threshold,
|
|
void * workspace,
|
|
void* output)
|
|
{
|
|
size_t nmsWorkspaceSize = _proposalsForwardNMSWorkspaceSize(batch_size, anc_size_num * anc_ratio_num,
|
|
rpn_height, rpn_width, MAX_BOX_NUM);
|
|
void* nmsWorkspace = workspace;
|
|
size_t proposalsSize = _proposalsForwardBboxWorkspaceSize(batch_size, anc_size_num * anc_ratio_num,
|
|
rpn_height, rpn_width);
|
|
const DType_t t_proposals = nvinfer1::DataType::kFLOAT;
|
|
const DLayout_t l_proposals = NC4HW;
|
|
void* proposals = nextWorkspacePtr((int8_t*) nmsWorkspace, nmsWorkspaceSize);
|
|
void* fg_scores = nextWorkspacePtr((int8_t*) proposals, proposalsSize);
|
|
size_t fg_scores_size = _proposalForwardFgScoresWorkspaceSize(batch_size,
|
|
anc_size_num * anc_ratio_num, rpn_height, rpn_width);
|
|
void* anchor_size_buf = nextWorkspacePtr((int8_t*) fg_scores, fg_scores_size);
|
|
void* anchor_ratio_buf = static_cast<void*>(static_cast<float*>(anchor_size_buf) + anc_size_num);
|
|
frcnnStatus_t status;
|
|
PLUGIN_CHECK_CUDA(
|
|
_copy_anchors_to_gpu(stream, ANCHOR_SIZES, anc_size_num, ANCHOR_RATIOS, anc_ratio_num, anchor_size_buf));
|
|
|
|
status = extractFgScores_gpu<float>(
|
|
stream, batch_size, anc_size_num * anc_ratio_num, rpn_height, rpn_width, rpn_prob, fg_scores);
|
|
PLUGIN_ASSERT(status == STATUS_SUCCESS);
|
|
PLUGIN_CHECK_CUDA(
|
|
_inverse_transform_wrapper(static_cast<const float*>(rpn_prob), static_cast<const float*>(rpn_regr), batch_size,
|
|
input_height, input_width, rpn_height, rpn_width, rpn_std_scaling, rpn_stride,
|
|
static_cast<float*>(anchor_size_buf), anc_size_num, static_cast<float*>(anchor_ratio_buf), anc_ratio_num,
|
|
bbox_min_size, static_cast<float*>(fg_scores), static_cast<float*>(proposals), stream));
|
|
|
|
status = nms(stream, batch_size, anc_size_num * anc_ratio_num * rpn_height * rpn_width, RPN_PRE_NMS_TOP_N,
|
|
MAX_BOX_NUM, nms_iou_threshold, nvinfer1::DataType::kFLOAT, NCHW, fg_scores, t_proposals, l_proposals,
|
|
proposals, workspace, nvinfer1::DataType::kFLOAT, output);
|
|
|
|
PLUGIN_ASSERT(status == STATUS_SUCCESS);
|
|
|
|
PLUGIN_CHECK_CUDA(
|
|
_normalize_rois(static_cast<float*>(output), batch_size, MAX_BOX_NUM, input_width, input_height, stream));
|
|
|
|
return STATUS_SUCCESS;
|
|
}
|
|
} // namespace plugin
|
|
} // namespace nvinfer1
|