wehub-resource-sync
231 lines
8.5 KiB
Plaintext
231 lines
8.5 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 "common/kernels/kernel.h"
|
|
#include "common/plugin.h"
|
|
#include <cuda_fp16.h>
|
|
#include <array>
|
|
|
|
|
|
namespace nvinfer1
|
|
{
|
|
namespace plugin
|
|
{
|
|
inline __device__ __half minus_fb(const __half& a, const __half& b)
|
|
{
|
|
#if __CUDA_ARCH__ >= 530
|
|
return a - b;
|
|
#else
|
|
return __float2half(__half2float(a) - __half2float(b));
|
|
#endif
|
|
}
|
|
|
|
inline __device__ float minus_fb(const float & a, const float & b) {
|
|
return a - b;
|
|
}
|
|
|
|
template <typename T_BBOX, typename T_SCORE, unsigned nthds_per_cta>
|
|
__launch_bounds__(nthds_per_cta)
|
|
__global__ void gatherTopDetections_kernel(
|
|
const bool shareLocation,
|
|
const int numImages,
|
|
const int numPredsPerClass,
|
|
const int numClasses,
|
|
const int topK,
|
|
const int keepTopK,
|
|
const int* indices,
|
|
const T_SCORE* scores,
|
|
const T_BBOX* bboxData,
|
|
int* keepCount,
|
|
T_BBOX* topDetections,
|
|
const T_SCORE score_shift)
|
|
{
|
|
if (keepTopK > topK)
|
|
return;
|
|
for (int i = blockIdx.x * nthds_per_cta + threadIdx.x;
|
|
i < numImages * keepTopK;
|
|
i += gridDim.x * nthds_per_cta)
|
|
{
|
|
const int imgId = i / keepTopK;
|
|
const int detId = i % keepTopK;
|
|
const int offset = imgId * numClasses * topK;
|
|
const int index = indices[offset + detId];
|
|
const T_SCORE score = scores[offset + detId];
|
|
/*
|
|
* It is also likely that there is "bad bounding boxes" in the keepTopK bounding boxes.
|
|
* We set the bounding boxes parameters as the parameters shown below.
|
|
* These data will only show up at the end of keepTopK bounding boxes since the bounding boxes were sorted previously.
|
|
* It is also not going to affect the count of valid bounding boxes (keepCount).
|
|
* These data will probably never be used (because we have keepCount).
|
|
*/
|
|
if (index == -1)
|
|
{
|
|
topDetections[i * 7] = imgId; // image id
|
|
topDetections[i * 7 + 1] = -1; // label
|
|
topDetections[i * 7 + 2] = 0; // confidence score
|
|
// score==0 will not pass the VisualizeBBox check
|
|
topDetections[i * 7 + 3] = 0; // bbox xmin
|
|
topDetections[i * 7 + 4] = 0; // bbox ymin
|
|
topDetections[i * 7 + 5] = 0; // bbox xmax
|
|
topDetections[i * 7 + 6] = 0; // bbox ymax
|
|
}
|
|
else
|
|
{
|
|
const int bboxOffset = imgId * (shareLocation ? numPredsPerClass : (numClasses * numPredsPerClass));
|
|
const int bboxId = ((shareLocation ? (index % numPredsPerClass)
|
|
: index % (numClasses * numPredsPerClass)) + bboxOffset) * 4;
|
|
topDetections[i * 7] = imgId; // image id
|
|
topDetections[i * 7 + 1] = (index % (numClasses * numPredsPerClass)) / numPredsPerClass; // label
|
|
topDetections[i * 7 + 2] = score; // confidence score
|
|
// subtract 1.0 score shift we added in sortScorePerClass
|
|
topDetections[i * 7 + 2] = minus_fb(topDetections[i * 7 + 2], score_shift);
|
|
const T_BBOX xMin = bboxData[bboxId];
|
|
const T_BBOX yMin = bboxData[bboxId + 1];
|
|
const T_BBOX xMax = bboxData[bboxId + 2];
|
|
const T_BBOX yMax = bboxData[bboxId + 3];
|
|
// clipped bbox xmin
|
|
topDetections[i * 7 + 3] = __saturatef(xMin);
|
|
// clipped bbox ymin
|
|
topDetections[i * 7 + 4] = __saturatef(yMin);
|
|
// clipped bbox xmax
|
|
topDetections[i * 7 + 5] = __saturatef(xMax);
|
|
// clipped bbox ymax
|
|
topDetections[i * 7 + 6] = __saturatef(yMax);
|
|
// Atomic add to increase the count of valid keepTopK bounding boxes
|
|
// Without having to do manual sync.
|
|
atomicAdd(&keepCount[i / keepTopK], 1);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T_BBOX, typename T_SCORE>
|
|
pluginStatus_t gatherTopDetections_gpu(
|
|
cudaStream_t stream,
|
|
const bool shareLocation,
|
|
const int numImages,
|
|
const int numPredsPerClass,
|
|
const int numClasses,
|
|
const int topK,
|
|
const int keepTopK,
|
|
const void* indices,
|
|
const void* scores,
|
|
const void* bboxData,
|
|
void* keepCount,
|
|
void* topDetections,
|
|
const float score_shift
|
|
)
|
|
{
|
|
CSC(cudaMemsetAsync(keepCount, 0, numImages * sizeof(int), stream), STATUS_FAILURE);
|
|
const int BS = 32;
|
|
const int GS = 32;
|
|
gatherTopDetections_kernel<T_BBOX, T_SCORE, BS><<<GS, BS, 0, stream>>>(shareLocation, numImages, numPredsPerClass,
|
|
numClasses, topK, keepTopK,
|
|
(int*) indices, (T_SCORE*) scores, (T_BBOX*) bboxData,
|
|
(int*) keepCount, (T_BBOX*) topDetections,
|
|
T_SCORE(score_shift));
|
|
|
|
CSC(cudaGetLastError(), STATUS_FAILURE);
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
// gatherTopDetections LAUNCH CONFIG
|
|
typedef pluginStatus_t (*gtdFunc)(cudaStream_t,
|
|
const bool,
|
|
const int,
|
|
const int,
|
|
const int,
|
|
const int,
|
|
const int,
|
|
const void*,
|
|
const void*,
|
|
const void*,
|
|
void*,
|
|
void*,
|
|
const float);
|
|
struct gtdLaunchConfig
|
|
{
|
|
DataType t_bbox;
|
|
DataType t_score;
|
|
gtdFunc function;
|
|
|
|
gtdLaunchConfig(DataType t_bbox, DataType t_score)
|
|
: t_bbox(t_bbox)
|
|
, t_score(t_score)
|
|
, function(nullptr)
|
|
{
|
|
}
|
|
gtdLaunchConfig(DataType t_bbox, DataType t_score, gtdFunc function)
|
|
: t_bbox(t_bbox)
|
|
, t_score(t_score)
|
|
, function(function)
|
|
{
|
|
}
|
|
bool operator==(gtdLaunchConfig const& other) const
|
|
{
|
|
return t_bbox == other.t_bbox && t_score == other.t_score;
|
|
}
|
|
};
|
|
|
|
using nvinfer1::DataType;
|
|
|
|
static std::array<gtdLaunchConfig, 2> gtdLCOptions = {
|
|
gtdLaunchConfig(DataType::kFLOAT, DataType::kFLOAT, gatherTopDetections_gpu<float, float>),
|
|
gtdLaunchConfig(DataType::kHALF, DataType::kHALF, gatherTopDetections_gpu<__half, __half>)
|
|
};
|
|
|
|
pluginStatus_t gatherTopDetections(
|
|
cudaStream_t stream,
|
|
const bool shareLocation,
|
|
const int numImages,
|
|
const int numPredsPerClass,
|
|
const int numClasses,
|
|
const int topK,
|
|
const int keepTopK,
|
|
const DataType DT_BBOX,
|
|
const DataType DT_SCORE,
|
|
const void* indices,
|
|
const void* scores,
|
|
const void* bboxData,
|
|
void* keepCount,
|
|
void* topDetections,
|
|
const float score_shift)
|
|
{
|
|
gtdLaunchConfig lc = gtdLaunchConfig(DT_BBOX, DT_SCORE);
|
|
for (unsigned i = 0; i < gtdLCOptions.size(); ++i)
|
|
{
|
|
if (lc == gtdLCOptions[i])
|
|
{
|
|
DEBUG_PRINTF("gatherTopDetections kernel %d\n", i);
|
|
return gtdLCOptions[i].function(stream,
|
|
shareLocation,
|
|
numImages,
|
|
numPredsPerClass,
|
|
numClasses,
|
|
topK,
|
|
keepTopK,
|
|
indices,
|
|
scores,
|
|
bboxData,
|
|
keepCount,
|
|
topDetections,
|
|
score_shift);
|
|
}
|
|
}
|
|
return STATUS_BAD_PARAM;
|
|
}
|
|
} // namespace plugin
|
|
} // namespace nvinfer1
|