Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add configurable input size for TLT MaskRCNN Plugin #986

Merged
merged 1 commit into from
Dec 29, 2020
Merged

Add configurable input size for TLT MaskRCNN Plugin #986

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 26 additions & 14 deletions plugin/generateDetectionPlugin/generateDetectionPlugin.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#include "generateDetectionPlugin.h"
#include "plugin.h"
#include <cuda_runtime_api.h>
#include <algorithm>

using namespace nvinfer1;
using namespace plugin;
Expand All @@ -40,6 +41,7 @@ GenerateDetectionPluginCreator::GenerateDetectionPluginCreator()
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();
Expand All @@ -62,6 +64,7 @@ const PluginFieldCollection* GenerateDetectionPluginCreator::getFieldNames()

IPluginV2Ext* GenerateDetectionPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
{
auto image_size = TLTMaskRCNNConfig::IMAGE_SHAPE;
const PluginField* fields = fc->fields;
for (int i = 0; i < fc->nbFields; ++i)
{
Expand All @@ -86,20 +89,27 @@ IPluginV2Ext* GenerateDetectionPluginCreator::createPlugin(const char* name, con
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);
return new GenerateDetection(mNbClasses, mKeepTopK, mScoreThreshold, mIOUThreshold, image_size);
};

IPluginV2Ext* GenerateDetectionPluginCreator::deserializePlugin(const char* name, const void* data, size_t length)
{
return new GenerateDetection(data, length);
};

GenerateDetection::GenerateDetection(int num_classes, int keep_topk, float score_threshold, float iou_threshold)
GenerateDetection::GenerateDetection(int num_classes, int keep_topk, float score_threshold, float iou_threshold, const nvinfer1::Dims& image_size)
: mNbClasses(num_classes)
, mKeepTopK(keep_topk)
, mScoreThreshold(score_threshold)
, mIOUThreshold(iou_threshold)
, mImageSize(image_size)
{
mBackgroundLabel = 0;
assert(mNbClasses > 0);
Expand Down Expand Up @@ -178,7 +188,7 @@ const char* GenerateDetection::getPluginNamespace() const

size_t GenerateDetection::getSerializationSize() const
{
return sizeof(int) * 2 + sizeof(float) * 2 + sizeof(int) * 2;
return sizeof(int) * 2 + sizeof(float) * 2 + sizeof(int) * 2 + sizeof(nvinfer1::Dims);
};

void GenerateDetection::serialize(void* buffer) const
Expand All @@ -190,6 +200,7 @@ void GenerateDetection::serialize(void* buffer) const
write(d, mIOUThreshold);
write(d, mMaxBatchSize);
write(d, mAnchorsCnt);
write(d, mImageSize);
ASSERT(d == a + getSerializationSize());
};

Expand All @@ -202,6 +213,7 @@ GenerateDetection::GenerateDetection(const void* data, size_t length)
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;
Expand Down Expand Up @@ -264,17 +276,17 @@ int GenerateDetection::enqueue(

// 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>(TLTMaskRCNNConfig::IMAGE_SHAPE.d[1]), // Image Height
static_cast<float>(TLTMaskRCNNConfig::IMAGE_SHAPE.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);
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;
Expand Down
4 changes: 3 additions & 1 deletion plugin/generateDetectionPlugin/generateDetectionPlugin.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ namespace plugin
class GenerateDetection : public IPluginV2Ext
{
public:
GenerateDetection(int num_classes, int keep_topk, float score_threshold, float iou_threshold);
GenerateDetection(int num_classes, int keep_topk, float score_threshold, float iou_threshold, const nvinfer1::Dims& image_size);

GenerateDetection(const void* data, size_t length);

Expand Down Expand Up @@ -103,6 +103,8 @@ class GenerateDetection : public IPluginV2Ext
RefineNMSParameters mParam;
std::shared_ptr<CudaBind<float>> mRegWeightDevice;

nvinfer1::Dims mImageSize;

std::string mNameSpace;
};

Expand Down
21 changes: 15 additions & 6 deletions plugin/multilevelCropAndResizePlugin/multilevelCropAndResizePlugin.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#include "multilevelCropAndResizePlugin.h"
#include "plugin.h"
#include <cuda_runtime_api.h>
#include <algorithm>

#include <fstream>

Expand All @@ -36,6 +37,7 @@ std::vector<PluginField> MultilevelCropAndResizePluginCreator::mPluginAttributes
MultilevelCropAndResizePluginCreator::MultilevelCropAndResizePluginCreator()
{
mPluginAttributes.emplace_back(PluginField("pooled_size", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("image_size", nullptr, PluginFieldType::kINT32, 3));

mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
Expand All @@ -58,6 +60,7 @@ const PluginFieldCollection* MultilevelCropAndResizePluginCreator::getFieldNames

IPluginV2Ext* MultilevelCropAndResizePluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
{
auto image_size = TLTMaskRCNNConfig::IMAGE_SHAPE;
const PluginField* fields = fc->fields;
for (int i = 0; i < fc->nbFields; ++i)
{
Expand All @@ -67,25 +70,31 @@ IPluginV2Ext* MultilevelCropAndResizePluginCreator::createPlugin(const char* nam
assert(fields[i].type == PluginFieldType::kINT32);
mPooledSize = *(static_cast<const int*>(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 MultilevelCropAndResize(mPooledSize);
return new MultilevelCropAndResize(mPooledSize, image_size);
};

IPluginV2Ext* MultilevelCropAndResizePluginCreator::deserializePlugin(const char* name, const void* data, size_t length)
{
return new MultilevelCropAndResize(data, length);
};

MultilevelCropAndResize::MultilevelCropAndResize(int pooled_size)
MultilevelCropAndResize::MultilevelCropAndResize(int pooled_size, const nvinfer1::Dims& image_size)
: mPooledSize({pooled_size, pooled_size})
{

assert(pooled_size > 0);
// shape
mInputHeight = TLTMaskRCNNConfig::IMAGE_SHAPE.d[1];
mInputWidth = TLTMaskRCNNConfig::IMAGE_SHAPE.d[2];
//Threshold to P3: Smaller -> P2
mThresh = (224*224) / (4.0f);
mInputHeight = image_size.d[1];
mInputWidth = image_size.d[2];
// Threshold to P3: Smaller -> P2
mThresh = (224 * 224) / (4.0f);
};

int MultilevelCropAndResize::getNbOutputs() const
Expand Down
2 changes: 1 addition & 1 deletion plugin/multilevelCropAndResizePlugin/multilevelCropAndResizePlugin.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ namespace plugin
class MultilevelCropAndResize : public IPluginV2Ext
{
public:
MultilevelCropAndResize(int pooled_size);
MultilevelCropAndResize(int pooled_size, const nvinfer1::Dims& image_size);

MultilevelCropAndResize(const void* data, size_t length);

Expand Down
55 changes: 31 additions & 24 deletions plugin/multilevelProposeROI/multilevelProposeROIPlugin.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
#include "plugin.h"
#include <cuda_runtime_api.h>
#include <iostream>
#include <algorithm>
#include <math.h>

#include <fstream>
Expand All @@ -43,6 +44,7 @@ MultilevelProposeROIPluginCreator::MultilevelProposeROIPluginCreator()
mPluginAttributes.emplace_back(PluginField("keep_topk", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("fg_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();
Expand All @@ -65,6 +67,7 @@ const PluginFieldCollection* MultilevelProposeROIPluginCreator::getFieldNames()

IPluginV2Ext* MultilevelProposeROIPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
{
auto image_size = TLTMaskRCNNConfig::IMAGE_SHAPE;
const PluginField* fields = fc->fields;
for (int i = 0; i < fc->nbFields; ++i)
{
Expand All @@ -89,20 +92,27 @@ IPluginV2Ext* MultilevelProposeROIPluginCreator::createPlugin(const char* name,
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 MultilevelProposeROI(mPreNMSTopK, mKeepTopK, mFGThreshold, mIOUThreshold);
return new MultilevelProposeROI(mPreNMSTopK, mKeepTopK, mFGThreshold, mIOUThreshold, image_size);
};

IPluginV2Ext* MultilevelProposeROIPluginCreator::deserializePlugin(const char* name, const void* data, size_t length)
{
return new MultilevelProposeROI(data, length);
};

MultilevelProposeROI::MultilevelProposeROI(int prenms_topk, int keep_topk, float fg_threshold, float iou_threshold)
MultilevelProposeROI::MultilevelProposeROI(int prenms_topk, int keep_topk, float fg_threshold, float iou_threshold, const nvinfer1::Dims image_size)
: mPreNMSTopK(prenms_topk)
, mKeepTopK(keep_topk)
, mFGThreshold(fg_threshold)
, mIOUThreshold(iou_threshold)
, mImageSize(image_size)
{
mBackgroundLabel = -1;
assert(mPreNMSTopK > 0);
Expand All @@ -121,7 +131,7 @@ MultilevelProposeROI::MultilevelProposeROI(int prenms_topk, int keep_topk, float

mFeatureCnt = TLTMaskRCNNConfig::MAX_LEVEL - TLTMaskRCNNConfig::MIN_LEVEL + 1;

generate_pyramid_anchors();
generate_pyramid_anchors(mImageSize);
};

int MultilevelProposeROI::getNbOutputs() const
Expand Down Expand Up @@ -224,7 +234,7 @@ const char* MultilevelProposeROI::getPluginNamespace() const

size_t MultilevelProposeROI::getSerializationSize() const
{
return sizeof(int) * 2 + sizeof(float) * 2 + sizeof(int) * (mFeatureCnt + 1);
return sizeof(int) * 2 + sizeof(float) * 2 + sizeof(int) * (mFeatureCnt + 1) + sizeof(nvinfer1::Dims);
};

void MultilevelProposeROI::serialize(void* buffer) const
Expand All @@ -239,6 +249,7 @@ void MultilevelProposeROI::serialize(void* buffer) const
{
write(d, mAnchorsCnt[i]);
}
write(d, mImageSize);
ASSERT(d == a + getSerializationSize());
};

Expand All @@ -257,6 +268,7 @@ MultilevelProposeROI::MultilevelProposeROI(const void* data, size_t length)
{
mAnchorsCnt.push_back(read<int>(d));
}
mImageSize = read<nvinfer1::Dims3>(d);
ASSERT(d == a + length);

mBackgroundLabel = -1;
Expand All @@ -273,7 +285,7 @@ MultilevelProposeROI::MultilevelProposeROI(const void* data, size_t length)

mType = DataType::kFLOAT;

generate_pyramid_anchors();
generate_pyramid_anchors(mImageSize);
};

void MultilevelProposeROI::check_valid_inputs(const nvinfer1::Dims* inputs, int nbInputDims)
Expand Down Expand Up @@ -329,9 +341,9 @@ Dims MultilevelProposeROI::getOutputDimensions(int index, const Dims* inputs, in
return proposals;
}

void MultilevelProposeROI::generate_pyramid_anchors()
void MultilevelProposeROI::generate_pyramid_anchors(const nvinfer1::Dims& image_size)
{
const auto image_dims = TLTMaskRCNNConfig::IMAGE_SHAPE;
const auto image_dims = image_size;

const auto& anchor_scale = TLTMaskRCNNConfig::RPN_ANCHOR_SCALE;
const auto& min_level = TLTMaskRCNNConfig::MIN_LEVEL;
Expand Down Expand Up @@ -388,23 +400,18 @@ int MultilevelProposeROI::enqueue(
{

MultilevelProposeROIWorkSpace proposal_ws(batch_size, mAnchorsCnt[i], mPreNMSTopK, mParam, mType);
status = MultilevelPropose(stream,
batch_size,
mAnchorsCnt[i],
mPreNMSTopK,
static_cast<float*>(mRegWeightDevice->mPtr),
static_cast<float>(TLTMaskRCNNConfig::IMAGE_SHAPE.d[1]), //Input Height
static_cast<float>(TLTMaskRCNNConfig::IMAGE_SHAPE.d[2]),
DataType::kFLOAT, // mType,
mParam,
proposal_ws,
workspace + kernel_workspace_offset,
inputs[2*i + 1], // inputs[object_score],
inputs[2*i], // inputs[bbox_delta]
mValidCnt->mPtr,
mAnchorBoxesDevice[i]->mPtr, // inputs[anchors]
mTempScores[i]->mPtr, //temp scores [batch_size, topk, 1]
mTempBboxes[i]->mPtr); //temp
status = MultilevelPropose(stream, batch_size, mAnchorsCnt[i], mPreNMSTopK,
static_cast<float*>(mRegWeightDevice->mPtr),
static_cast<float>(mImageSize.d[1]), // Input Height
static_cast<float>(mImageSize.d[2]),
DataType::kFLOAT, // mType,
mParam, proposal_ws, workspace + kernel_workspace_offset,
inputs[2 * i + 1], // inputs[object_score],
inputs[2 * i], // inputs[bbox_delta]
mValidCnt->mPtr,
mAnchorBoxesDevice[i]->mPtr, // inputs[anchors]
mTempScores[i]->mPtr, // temp scores [batch_size, topk, 1]
mTempBboxes[i]->mPtr); // temp
assert(status == cudaSuccess);
kernel_workspace_offset += proposal_ws.totalSize;
}
Expand Down
5 changes: 3 additions & 2 deletions plugin/multilevelProposeROI/multilevelProposeROIPlugin.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ namespace plugin
class MultilevelProposeROI : public IPluginV2Ext
{
public:
MultilevelProposeROI(int prenms_topk, int keep_topk, float fg_threshold, float iou_threshold);
MultilevelProposeROI(int prenms_topk, int keep_topk, float fg_threshold, float iou_threshold, const nvinfer1::Dims image_size);

MultilevelProposeROI(const void* data, size_t length);

Expand Down Expand Up @@ -88,7 +88,7 @@ class MultilevelProposeROI : public IPluginV2Ext

private:
void check_valid_inputs(const nvinfer1::Dims* inputs, int nbInputDims);
void generate_pyramid_anchors();
void generate_pyramid_anchors(const nvinfer1::Dims& image_size);

int mBackgroundLabel;
int mPreNMSTopK;
Expand All @@ -111,6 +111,7 @@ class MultilevelProposeROI : public IPluginV2Ext
float** mDeviceBboxes;
std::shared_ptr<CudaBind<float>> mRegWeightDevice;

nvinfer1::Dims mImageSize;
nvinfer1::DataType mType;
RefineNMSParameters mParam;

Expand Down