Rajeev Rao
337 lines
11 KiB
C++
337 lines
11 KiB
C++
/*
|
|
* Copyright (c) 2021, NVIDIA CORPORATION. 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 "generateDetectionPlugin.h"
|
|
#include "plugin.h"
|
|
#include <algorithm>
|
|
#include <cuda_runtime_api.h>
|
|
|
|
using namespace nvinfer1;
|
|
using namespace plugin;
|
|
using nvinfer1::plugin::GenerateDetection;
|
|
using nvinfer1::plugin::GenerateDetectionPluginCreator;
|
|
|
|
#include <fstream>
|
|
|
|
namespace
|
|
{
|
|
const char* GENERATEDETECTION_PLUGIN_VERSION{"1"};
|
|
const char* GENERATEDETECTION_PLUGIN_NAME{"GenerateDetection_TRT"};
|
|
} // namespace
|
|
|
|
PluginFieldCollection GenerateDetectionPluginCreator::mFC{};
|
|
std::vector<PluginField> GenerateDetectionPluginCreator::mPluginAttributes;
|
|
|
|
GenerateDetectionPluginCreator::GenerateDetectionPluginCreator() noexcept
|
|
{
|
|
|
|
mPluginAttributes.emplace_back(PluginField("num_classes", nullptr, PluginFieldType::kINT32, 1));
|
|
mPluginAttributes.emplace_back(PluginField("keep_topk", nullptr, PluginFieldType::kINT32, 1));
|
|
mPluginAttributes.emplace_back(PluginField("score_threshold", nullptr, PluginFieldType::kFLOAT32, 1));
|
|
mPluginAttributes.emplace_back(PluginField("iou_threshold", nullptr, PluginFieldType::kFLOAT32, 1));
|
|
mPluginAttributes.emplace_back(PluginField("image_size", nullptr, PluginFieldType::kINT32, 3));
|
|
|
|
mFC.nbFields = mPluginAttributes.size();
|
|
mFC.fields = mPluginAttributes.data();
|
|
}
|
|
|
|
const char* GenerateDetectionPluginCreator::getPluginName() const noexcept
|
|
{
|
|
return GENERATEDETECTION_PLUGIN_NAME;
|
|
};
|
|
|
|
const char* GenerateDetectionPluginCreator::getPluginVersion() const noexcept
|
|
{
|
|
return GENERATEDETECTION_PLUGIN_VERSION;
|
|
};
|
|
|
|
const PluginFieldCollection* GenerateDetectionPluginCreator::getFieldNames() noexcept
|
|
{
|
|
return &mFC;
|
|
};
|
|
|
|
IPluginV2Ext* GenerateDetectionPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) noexcept
|
|
{
|
|
auto image_size = TLTMaskRCNNConfig::IMAGE_SHAPE;
|
|
const PluginField* fields = fc->fields;
|
|
for (int i = 0; i < fc->nbFields; ++i)
|
|
{
|
|
const char* attrName = fields[i].name;
|
|
if (!strcmp(attrName, "num_classes"))
|
|
{
|
|
assert(fields[i].type == PluginFieldType::kINT32);
|
|
mNbClasses = *(static_cast<const int*>(fields[i].data));
|
|
}
|
|
if (!strcmp(attrName, "keep_topk"))
|
|
{
|
|
assert(fields[i].type == PluginFieldType::kINT32);
|
|
mKeepTopK = *(static_cast<const int*>(fields[i].data));
|
|
}
|
|
if (!strcmp(attrName, "score_threshold"))
|
|
{
|
|
assert(fields[i].type == PluginFieldType::kFLOAT32);
|
|
mScoreThreshold = *(static_cast<const float*>(fields[i].data));
|
|
}
|
|
if (!strcmp(attrName, "iou_threshold"))
|
|
{
|
|
assert(fields[i].type == PluginFieldType::kFLOAT32);
|
|
mIOUThreshold = *(static_cast<const float*>(fields[i].data));
|
|
}
|
|
if (!strcmp(attrName, "image_size"))
|
|
{
|
|
assert(fields[i].type == PluginFieldType::kINT32);
|
|
const auto dims = static_cast<const int32_t*>(fields[i].data);
|
|
std::copy_n(dims, 3, image_size.d);
|
|
}
|
|
}
|
|
return new GenerateDetection(mNbClasses, mKeepTopK, mScoreThreshold, mIOUThreshold, image_size);
|
|
};
|
|
|
|
IPluginV2Ext* GenerateDetectionPluginCreator::deserializePlugin(const char* name, const void* data, size_t length) noexcept
|
|
{
|
|
return new GenerateDetection(data, length);
|
|
};
|
|
|
|
GenerateDetection::GenerateDetection(
|
|
int num_classes, int keep_topk, float score_threshold, float iou_threshold, const nvinfer1::Dims& image_size) noexcept
|
|
: mNbClasses(num_classes)
|
|
, mKeepTopK(keep_topk)
|
|
, mScoreThreshold(score_threshold)
|
|
, mIOUThreshold(iou_threshold)
|
|
, mImageSize(image_size)
|
|
{
|
|
mBackgroundLabel = 0;
|
|
assert(mNbClasses > 0);
|
|
assert(mKeepTopK > 0);
|
|
assert(score_threshold >= 0.0f);
|
|
assert(iou_threshold > 0.0f);
|
|
|
|
mParam.backgroundLabelId = 0;
|
|
mParam.numClasses = mNbClasses;
|
|
mParam.keepTopK = mKeepTopK;
|
|
mParam.scoreThreshold = mScoreThreshold;
|
|
mParam.iouThreshold = mIOUThreshold;
|
|
|
|
mType = DataType::kFLOAT;
|
|
};
|
|
|
|
int GenerateDetection::getNbOutputs() const noexcept
|
|
{
|
|
return 1;
|
|
};
|
|
|
|
int GenerateDetection::initialize() noexcept
|
|
{
|
|
// Init the regWeight [10, 10, 5, 5]
|
|
mRegWeightDevice = std::make_shared<CudaBind<float>>(4);
|
|
CUASSERT(cudaMemcpy(static_cast<void*>(mRegWeightDevice->mPtr),
|
|
static_cast<const void*>(TLTMaskRCNNConfig::DETECTION_REG_WEIGHTS), sizeof(float) * 4, cudaMemcpyHostToDevice));
|
|
|
|
//@Init the mValidCnt and mDecodedBboxes for max batch size
|
|
std::vector<int> tempValidCnt(mMaxBatchSize, mAnchorsCnt);
|
|
|
|
mValidCnt = std::make_shared<CudaBind<int>>(mMaxBatchSize);
|
|
|
|
CUASSERT(cudaMemcpy(
|
|
mValidCnt->mPtr, static_cast<void*>(tempValidCnt.data()), sizeof(int) * mMaxBatchSize, cudaMemcpyHostToDevice));
|
|
|
|
return 0;
|
|
};
|
|
|
|
void GenerateDetection::terminate() noexcept {};
|
|
|
|
void GenerateDetection::destroy() noexcept
|
|
{
|
|
delete this;
|
|
};
|
|
|
|
bool GenerateDetection::supportsFormat(DataType type, PluginFormat format) const noexcept
|
|
{
|
|
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
|
|
};
|
|
|
|
const char* GenerateDetection::getPluginType() const noexcept
|
|
{
|
|
return "GenerateDetection_TRT";
|
|
};
|
|
|
|
const char* GenerateDetection::getPluginVersion() const noexcept
|
|
{
|
|
return "1";
|
|
};
|
|
|
|
IPluginV2Ext* GenerateDetection::clone() const noexcept
|
|
{
|
|
return new GenerateDetection(*this);
|
|
};
|
|
|
|
void GenerateDetection::setPluginNamespace(const char* libNamespace) noexcept
|
|
{
|
|
mNameSpace = libNamespace;
|
|
};
|
|
|
|
const char* GenerateDetection::getPluginNamespace() const noexcept
|
|
{
|
|
return mNameSpace.c_str();
|
|
}
|
|
|
|
size_t GenerateDetection::getSerializationSize() const noexcept
|
|
{
|
|
return sizeof(int) * 2 + sizeof(float) * 2 + sizeof(int) * 2 + sizeof(nvinfer1::Dims);
|
|
};
|
|
|
|
void GenerateDetection::serialize(void* buffer) const noexcept
|
|
{
|
|
char *d = reinterpret_cast<char*>(buffer), *a = d;
|
|
write(d, mNbClasses);
|
|
write(d, mKeepTopK);
|
|
write(d, mScoreThreshold);
|
|
write(d, mIOUThreshold);
|
|
write(d, mMaxBatchSize);
|
|
write(d, mAnchorsCnt);
|
|
write(d, mImageSize);
|
|
ASSERT(d == a + getSerializationSize());
|
|
};
|
|
|
|
GenerateDetection::GenerateDetection(const void* data, size_t length) noexcept
|
|
{
|
|
const char *d = reinterpret_cast<const char*>(data), *a = d;
|
|
int num_classes = read<int>(d);
|
|
int keep_topk = read<int>(d);
|
|
float score_threshold = read<float>(d);
|
|
float iou_threshold = read<float>(d);
|
|
mMaxBatchSize = read<int>(d);
|
|
mAnchorsCnt = read<int>(d);
|
|
mImageSize = read<nvinfer1::Dims3>(d);
|
|
ASSERT(d == a + length);
|
|
|
|
mNbClasses = num_classes;
|
|
mKeepTopK = keep_topk;
|
|
mScoreThreshold = score_threshold;
|
|
mIOUThreshold = iou_threshold;
|
|
|
|
mParam.backgroundLabelId = 0;
|
|
mParam.numClasses = mNbClasses;
|
|
mParam.keepTopK = mKeepTopK;
|
|
mParam.scoreThreshold = mScoreThreshold;
|
|
mParam.iouThreshold = mIOUThreshold;
|
|
|
|
mType = DataType::kFLOAT;
|
|
};
|
|
|
|
void GenerateDetection::check_valid_inputs(const nvinfer1::Dims* inputs, int nbInputDims) noexcept
|
|
{
|
|
// classifier_delta_bbox[N, anchors, num_classes*4, 1, 1]
|
|
// classifier_class[N, anchors, num_classes, 1, 1]
|
|
// rpn_rois[N, anchors, 4]
|
|
assert(nbInputDims == 3);
|
|
|
|
// score
|
|
assert(inputs[1].nbDims == 4 && inputs[1].d[1] == mNbClasses);
|
|
// delta_bbox
|
|
assert(inputs[0].nbDims == 4 && inputs[0].d[1] == mNbClasses * 4);
|
|
// roi
|
|
assert(inputs[2].nbDims == 2 && inputs[2].d[1] == 4);
|
|
};
|
|
|
|
size_t GenerateDetection::getWorkspaceSize(int batch_size) const noexcept
|
|
{
|
|
RefineDetectionWorkSpace refine(batch_size, mAnchorsCnt, mParam, mType);
|
|
return refine.totalSize;
|
|
};
|
|
|
|
Dims GenerateDetection::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) noexcept
|
|
{
|
|
|
|
check_valid_inputs(inputs, nbInputDims);
|
|
assert(index == 0);
|
|
|
|
// [N, anchors, (y1, x1, y2, x2, class_id, score)]
|
|
nvinfer1::Dims detections;
|
|
|
|
detections.nbDims = 2;
|
|
// number of anchors
|
|
detections.d[0] = mKeepTopK;
|
|
detections.d[1] = 6;
|
|
|
|
return detections;
|
|
}
|
|
|
|
int32_t GenerateDetection::enqueue(
|
|
int32_t batch_size, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
|
|
{
|
|
|
|
void* detections = outputs[0];
|
|
|
|
// refine detection
|
|
RefineDetectionWorkSpace refDetcWorkspace(batch_size, mAnchorsCnt, mParam, mType);
|
|
cudaError_t status
|
|
= DetectionPostProcess(stream, batch_size, mAnchorsCnt, static_cast<float*>(mRegWeightDevice->mPtr),
|
|
static_cast<float>(mImageSize.d[1]), // Image Height
|
|
static_cast<float>(mImageSize.d[2]), // Image Width
|
|
DataType::kFLOAT, // mType,
|
|
mParam, refDetcWorkspace, workspace,
|
|
inputs[1], // inputs[InScore]
|
|
inputs[0], // inputs[InDelta],
|
|
mValidCnt->mPtr, // inputs[InCountValid],
|
|
inputs[2], // inputs[ROI]
|
|
detections);
|
|
|
|
assert(status == cudaSuccess);
|
|
return status;
|
|
};
|
|
|
|
DataType GenerateDetection::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept
|
|
{
|
|
// Only DataType::kFLOAT is acceptable by the plugin layer
|
|
return DataType::kFLOAT;
|
|
}
|
|
|
|
// Return true if output tensor is broadcast across a batch.
|
|
bool GenerateDetection::isOutputBroadcastAcrossBatch(
|
|
int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const noexcept
|
|
{
|
|
return false;
|
|
}
|
|
|
|
// Return true if plugin can use input that is broadcast across batch without replication.
|
|
bool GenerateDetection::canBroadcastInputAcrossBatch(int inputIndex) const noexcept
|
|
{
|
|
return false;
|
|
}
|
|
|
|
// Configure the layer with input and output data types.
|
|
void GenerateDetection::configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,
|
|
const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast,
|
|
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) noexcept
|
|
{
|
|
check_valid_inputs(inputDims, nbInputs);
|
|
assert(inputDims[0].d[0] == inputDims[1].d[0] && inputDims[1].d[0] == inputDims[2].d[0]);
|
|
|
|
mAnchorsCnt = inputDims[2].d[0];
|
|
mType = inputTypes[0];
|
|
mMaxBatchSize = maxBatchSize;
|
|
}
|
|
|
|
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
|
|
void GenerateDetection::attachToContext(
|
|
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept
|
|
{
|
|
}
|
|
|
|
// Detach the plugin object from its execution context.
|
|
void GenerateDetection::detachFromContext() noexcept {}
|