Skip to content

Commit

Permalink
Merge pull request #679 from rajeevsrao/dev/gridAnchorRect
Browse files Browse the repository at this point in the history 
Re-enable GridAnchorRect_TRT plugin with rectangular feature maps
  • Loading branch information
@rajeevsrao
rajeevsrao authored Jul 12, 2020
2 parents 9ee8450 + c9b4e4d commit ead9e8b
Show file tree
Hide file tree
Showing 6 changed files with 129 additions and 124 deletions.
1 change: 1 addition & 0 deletions plugin/InferPlugin.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,7 @@ extern "C" {
bool initLibNvInferPlugins(void* logger, const char* libNamespace)
{
initializePlugin<nvinfer1::plugin::GridAnchorPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::GridAnchorRectPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::NMSPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::ReorgPluginCreator>(logger, libNamespace);
initializePlugin<nvinfer1::plugin::RegionPluginCreator>(logger, libNamespace);
Expand Down
1 change: 1 addition & 0 deletions plugin/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
| [geluPlugin](geluPlugin) | CustomGeluPluginDynamic | 1 |
| [generateDetectionPlugin](generateDetectionPlugin) | GenerateDetection_TRT | 1 |
| [gridAnchorPlugin](gridAnchorPlugin) | GridAnchor_TRT | 1 |
| [gridAnchorRectPlugin](gridAnchorPlugin) | GridAnchorRect_TRT | 1 |
| [groupNormalizationPlugin](groupNormalizationPlugin) | GroupNormalizationPlugin | 1 |
| [instanceNormalizationPlugin](instanceNormalizationPlugin) | InstanceNormalization_TRT | 1 |
| [multilevelCropAndResizePlugin](multilevelCropAndResizePlugin) | MultilevelCropAndResize_TRT | 1 |
Expand Down
153 changes: 59 additions & 94 deletions plugin/common/kernels/gridAnchorLayer.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,105 +19,73 @@

using nvinfer1::plugin::reduced_divisor;
template <unsigned nthdsPerCTA>
__launch_bounds__(nthdsPerCTA)
__global__ void gridAnchorKernel(
const GridAnchorParameters param,
const int numAspectRatios,
reduced_divisor divObj,
const float* widths,
const float* heights,
float* outputData
)
{
// output dims: (H, W, param.numMinSize, (1+haveMaxSize+numAR-1), 4)
const int dim = param.H * param.W * numAspectRatios;
/*
* Parameters used to calculate the bounding box coordinates back to input image scale
* Normally we calculate the anchorStride = image_input_size (in pixel) / feature_map_size
* Here we do not use image_input_size for the moment
* Instead we use 1.0
* The coordinates calculated are scaled by the input image size.
* Most of the coordinates will be in a range of [0, 1], except for the bounding box coordinates going outside of the image
* Every coordinate will go back to the pixel coordinates in the input image if being multiplied by image_input_size
* Here we implicitly assumes the image input and feature map are square
*/
float anchorStride = (1.0 / param.H);
float anchorOffset = 0.5 * anchorStride;
__launch_bounds__(nthdsPerCTA)
__global__ void gridAnchorKernel(
const GridAnchorParameters param,
const int numAspectRatios,
reduced_divisor divObj,
const float* widths,
const float* heights,
float* outputData
)
{
// output dims: (H, W, param.numMinSize, (1+haveMaxSize+numAR-1), 4)
const int dim = param.H * param.W * numAspectRatios;

int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= dim)
return;
int arId, currIndex;
divObj.divmod(tid, currIndex, arId);
/*
* Parameters used to calculate the bounding box coordinates back to input image scale
* Normally we calculate the anchorStride = image_input_size (in pixel) / feature_map_size
* Here we do not use image_input_size for the moment
* Instead we use 1.0
* The coordinates calculated are scaled by the input image size.
* Most of the coordinates will be in a range of [0, 1], except for the bounding box coordinates going outside of the image
* Every coordinate will go back to the pixel coordinates in the input image if being multiplied by image_input_size
* Here we implicitly assumes the image input and feature map are square
*/
float anchorStrideH = (1.0 / param.H);
float anchorStrideW = (1.0 / param.W);
float anchorOffsetH = 0.5 * anchorStrideH;
float anchorOffsetW = 0.5 * anchorStrideW;

const int w = currIndex % param.W;
const int h = currIndex / param.W;

// Center coordinates
float yC = h * anchorStride + anchorOffset;
float xC = w * anchorStride + anchorOffset;

// x_min, y_min
float xMin = xC - 0.5 * widths[arId];
float yMin = yC - 0.5 * heights[arId];
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= dim)
{
return;
}

// x_max, y_max
float xMax = xC + 0.5 * widths[arId];
float yMax = yC + 0.5 * heights[arId];
int arId, currIndex;
divObj.divmod(tid, currIndex, arId);

outputData[tid * 4] = xMin;
outputData[tid * 4 + 1] = yMin;
outputData[tid * 4 + 2] = xMax;
outputData[tid * 4 + 3] = yMax;
const int w = currIndex % param.W;
const int h = currIndex / param.W;

// Remember to move the output cursor
float* output = outputData + dim * 4;
// Center coordinates
float yC = h * anchorStrideH + anchorOffsetH;
float xC = w * anchorStrideW + anchorOffsetW;

// Simply copying the variance
output[tid * 4] = param.variance[0];
output[tid * 4 + 1] = param.variance[1];
output[tid * 4 + 2] = param.variance[2];
output[tid * 4 + 3] = param.variance[3];
// x_min, y_min
float xMin = xC - 0.5 * widths[arId];
float yMin = yC - 0.5 * heights[arId];

}
// x_max, y_max
float xMax = xC + 0.5 * widths[arId];
float yMax = yC + 0.5 * heights[arId];

pluginStatus_t anchorGridInference(
cudaStream_t stream,
const GridAnchorParameters param,
const int numAspectRatios,
const void* widths,
const void* heights,
void* outputData
)
{
const int dim = param.H * param.W * numAspectRatios;
reduced_divisor divObj(numAspectRatios);
if (dim > 5120)
{
const int BS = 128;
const int GS = (dim + BS - 1) / BS;
gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(param, numAspectRatios, divObj,
(const float*) widths, (const float*) heights,
(float*) outputData);
outputData[tid * 4] = xMin;
outputData[tid * 4 + 1] = yMin;
outputData[tid * 4 + 2] = xMax;
outputData[tid * 4 + 3] = yMax;

}
else
{
const int BS = 32;
const int GS = (dim + BS - 1) / BS;
gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(param, numAspectRatios, divObj,
(const float*) widths, (const float*) heights,
(float*) outputData);
// Remember to move the output cursor
float* output = outputData + dim * 4;

}
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
// Simply copying the variance
output[tid * 4] = param.variance[0];
output[tid * 4 + 1] = param.variance[1];
output[tid * 4 + 2] = param.variance[2];
output[tid * 4 + 3] = param.variance[3];
}

namespace nvinfer1
{
namespace plugin
{
pluginStatus_t anchorGridInference(
cudaStream_t stream,
const GridAnchorParameters param,
Expand All @@ -134,22 +102,19 @@ pluginStatus_t anchorGridInference(
const int BS = 128;
const int GS = (dim + BS - 1) / BS;
gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(param, numAspectRatios, divObj,
(const float*) widths, (const float*) heights,
(float*) outputData);
(const float*) widths, (const float*) heights,
(float*) outputData);

}
else
{
const int BS = 32;
const int GS = (dim + BS - 1) / BS;
gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(param, numAspectRatios, divObj,
(const float*) widths, (const float*) heights,
(float*) outputData);

(const float*) widths, (const float*) heights,
(float*) outputData);
}
CSC(cudaGetLastError(), STATUS_FAILURE);
return STATUS_SUCCESS;
}

}
}
65 changes: 41 additions & 24 deletions plugin/gridAnchorPlugin/gridAnchorPlugin.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,19 +22,18 @@
#include <vector>

using namespace nvinfer1;
using nvinfer1::plugin::GridAnchorGenerator;
using nvinfer1::plugin::GridAnchorPluginCreator;

namespace
{
const char* GRID_ANCHOR_PLUGIN_VERSION{"1"};
const char* GRID_ANCHOR_PLUGIN_NAME{"GridAnchor_TRT"};
std::string GRID_ANCHOR_PLUGIN_NAMES[] = {"GridAnchor_TRT", "GridAnchorRect_TRT"};
const char* GRID_ANCHOR_PLUGIN_VERSION = "1";
} // namespace
PluginFieldCollection GridAnchorPluginCreator::mFC{};
std::vector<PluginField> GridAnchorPluginCreator::mPluginAttributes;

GridAnchorGenerator::GridAnchorGenerator(const GridAnchorParameters* paramIn, int mNumLayers)
: mNumLayers(mNumLayers)
PluginFieldCollection GridAnchorBasePluginCreator::mFC{};
std::vector<PluginField> GridAnchorBasePluginCreator::mPluginAttributes;

GridAnchorGenerator::GridAnchorGenerator(const GridAnchorParameters* paramIn, int numLayers, const char *name)
: mNumLayers(numLayers), mPluginName(name)
{
CUASSERT(cudaMallocHost((void**) &mNumPriors, mNumLayers * sizeof(int)));
CUASSERT(cudaMallocHost((void**) &mDeviceWidths, mNumLayers * sizeof(Weights)));
Expand Down Expand Up @@ -121,7 +120,8 @@ GridAnchorGenerator::GridAnchorGenerator(const GridAnchorParameters* paramIn, in
}
}

GridAnchorGenerator::GridAnchorGenerator(const void* data, size_t length)
GridAnchorGenerator::GridAnchorGenerator(const void* data, size_t length, const char *name) :
mPluginName(name)
{
const char *d = reinterpret_cast<const char*>(data), *a = d;
mNumLayers = read<int>(d);
Expand Down Expand Up @@ -276,14 +276,15 @@ bool GridAnchorGenerator::supportsFormat(DataType type, PluginFormat format) con

const char* GridAnchorGenerator::getPluginType() const
{
return GRID_ANCHOR_PLUGIN_NAME;
return mPluginName.c_str();
}

const char* GridAnchorGenerator::getPluginVersion() const
{
return GRID_ANCHOR_PLUGIN_VERSION;
}


// Set plugin namespace
void GridAnchorGenerator::setPluginNamespace(const char* pluginNamespace)
{
Expand Down Expand Up @@ -341,12 +342,12 @@ void GridAnchorGenerator::destroy()

IPluginV2Ext* GridAnchorGenerator::clone() const
{
IPluginV2Ext* plugin = new GridAnchorGenerator(mParam.data(), mNumLayers);
IPluginV2Ext* plugin = new GridAnchorGenerator(mParam.data(), mNumLayers, mPluginName.c_str());
plugin->setPluginNamespace(mPluginNamespace.c_str());
return plugin;
}

GridAnchorPluginCreator::GridAnchorPluginCreator()
GridAnchorBasePluginCreator::GridAnchorBasePluginCreator()
{
mPluginAttributes.emplace_back(PluginField("minSize", nullptr, PluginFieldType::kFLOAT32, 1));
mPluginAttributes.emplace_back(PluginField("maxSize", nullptr, PluginFieldType::kFLOAT32, 1));
Expand All @@ -359,29 +360,31 @@ GridAnchorPluginCreator::GridAnchorPluginCreator()
mFC.fields = mPluginAttributes.data();
}

const char* GridAnchorPluginCreator::getPluginName() const
const char* GridAnchorBasePluginCreator::getPluginName() const
{
return GRID_ANCHOR_PLUGIN_NAME;
return mPluginName.c_str();
}

const char* GridAnchorPluginCreator::getPluginVersion() const
const char* GridAnchorBasePluginCreator::getPluginVersion() const
{
return GRID_ANCHOR_PLUGIN_VERSION;
}

const PluginFieldCollection* GridAnchorPluginCreator::getFieldNames()
const PluginFieldCollection* GridAnchorBasePluginCreator::getFieldNames()
{
return &mFC;
}

IPluginV2Ext* GridAnchorPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
IPluginV2Ext* GridAnchorBasePluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
{
float minScale = 0.2F, maxScale = 0.95F;
int numLayers = 6;
std::vector<float> aspectRatios;
std::vector<int> fMapShapes;
std::vector<float> layerVariances;
const PluginField* fields = fc->fields;

const bool isFMapRect = (GRID_ANCHOR_PLUGIN_NAMES[1] == mPluginName);
for (int i = 0; i < fc->nbFields; ++i)
{
const char* attrName = fields[i].name;
Expand Down Expand Up @@ -428,6 +431,7 @@ IPluginV2Ext* GridAnchorPluginCreator::createPlugin(const char* name, const Plug
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
int size = fields[i].length;
ASSERT(!isFMapRect || (size % 2 == 0));
fMapShapes.reserve(size);
const int* fMap = static_cast<const int*>(fields[i].data);
for (int j = 0; j < size; j++)
Expand All @@ -442,7 +446,8 @@ IPluginV2Ext* GridAnchorPluginCreator::createPlugin(const char* name, const Plug
std::vector<float> firstLayerAspectRatios;

ASSERT(numLayers > 0);
ASSERT((int) fMapShapes.size() == numLayers);
const int numExpectedLayers = static_cast<int>(fMapShapes.size()) >> (isFMapRect ? 1 : 0);
ASSERT(numExpectedLayers == numLayers);

int numFirstLayerARs = 3;
// First layer only has the first 3 aspect ratios from aspectRatios
Expand All @@ -457,30 +462,42 @@ IPluginV2Ext* GridAnchorPluginCreator::createPlugin(const char* name, const Plug
// One set of box parameters for one layer
for (int i = 0; i < numLayers; i++)
{
int hOffset = (isFMapRect ? i * 2 : i);
int wOffset = (isFMapRect ? i * 2 + 1 : i);
// Only the first layer is different
if (i == 0)
{
boxParams[i] = {minScale, maxScale, firstLayerAspectRatios.data(), (int) firstLayerAspectRatios.size(),
fMapShapes[i], fMapShapes[i],
fMapShapes[hOffset], fMapShapes[wOffset],
{layerVariances[0], layerVariances[1], layerVariances[2], layerVariances[3]}};
}
else
{
boxParams[i] = {minScale, maxScale, aspectRatios.data(), (int) aspectRatios.size(), fMapShapes[i],
fMapShapes[i], {layerVariances[0], layerVariances[1], layerVariances[2], layerVariances[3]}};
boxParams[i] = {minScale, maxScale, aspectRatios.data(), (int) aspectRatios.size(), fMapShapes[hOffset],
fMapShapes[wOffset], {layerVariances[0], layerVariances[1], layerVariances[2], layerVariances[3]}};
}
}

GridAnchorGenerator* obj = new GridAnchorGenerator(boxParams.data(), numLayers);
GridAnchorGenerator* obj = new GridAnchorGenerator(boxParams.data(), numLayers, mPluginName.c_str());
obj->setPluginNamespace(mNamespace.c_str());
return obj;
}

IPluginV2Ext* GridAnchorPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
IPluginV2Ext* GridAnchorBasePluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
{
// This object will be deleted when the network is destroyed, which will
// call GridAnchor::destroy()
GridAnchorGenerator* obj = new GridAnchorGenerator(serialData, serialLength);
GridAnchorGenerator* obj = new GridAnchorGenerator(serialData, serialLength, mPluginName.c_str());
obj->setPluginNamespace(mNamespace.c_str());
return obj;
}

GridAnchorPluginCreator::GridAnchorPluginCreator()
{
mPluginName = GRID_ANCHOR_PLUGIN_NAMES[0];
}

GridAnchorRectPluginCreator::GridAnchorRectPluginCreator()
{
mPluginName = GRID_ANCHOR_PLUGIN_NAMES[1];
}
Loading

0 comments on commit ead9e8b

Please sign in to comment.