Re-enable GridAnchorRect_TRT plugin with rectangular feature maps

plugin/InferPlugin.cpp

@@ -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);

plugin/README.md

@@ -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 |

plugin/common/kernels/gridAnchorLayer.cu

@@ -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,
@@ -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;
 }
 
-}
-}

plugin/gridAnchorPlugin/gridAnchorPlugin.cpp

@@ -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)));
@@ -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);
@@ -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)
 {
@@ -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));
@@ -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;
@@ -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++)
@@ -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
@@ -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];
+}