Method getOutputDimensions

plugin/efficientNMSPlugin/efficientNMSPlugin.cpp:166–234

}

DimsExprs EfficientNMSPlugin::getOutputDimensions(
    int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
{
    try
    {
        DimsExprs out_dim;

        // When pad per class is set, the output size may need to be reduced:
        // i.e.: outputBoxes = min(outputBoxes, outputBoxesPerClass * numClasses)
        // As the number of classes may not be static, numOutputBoxes must be a dynamic
        // expression. The corresponding parameter can not be set at this time, so the
        // value will be calculated again in configurePlugin() and the param overwritten.
        IDimensionExpr const* numOutputBoxes = exprBuilder.constant(mParam.numOutputBoxes);
        if (mParam.padOutputBoxesPerClass && mParam.numOutputBoxesPerClass > 0)
        {
            IDimensionExpr const* numOutputBoxesPerClass = exprBuilder.constant(mParam.numOutputBoxesPerClass);
            IDimensionExpr const* numClasses = inputs[1].d[2];
            numOutputBoxes = exprBuilder.operation(DimensionOperation::kMIN, *numOutputBoxes,
                *exprBuilder.operation(DimensionOperation::kPROD, *numOutputBoxesPerClass, *numClasses));
        }

        if (mParam.outputONNXIndices)
        {
            // ONNX NMS
            PLUGIN_ASSERT(outputIndex == 0);

            // detection_indices
            out_dim.nbDims = 2;
            out_dim.d[0] = exprBuilder.operation(DimensionOperation::kPROD, *inputs[0].d[0], *numOutputBoxes);
            out_dim.d[1] = exprBuilder.constant(3);
        }
        else
        {
            // Standard NMS
            PLUGIN_ASSERT(outputIndex >= 0 && outputIndex <= 3);

            // num_detections
            if (outputIndex == 0)
            {
                out_dim.nbDims = 2;
                out_dim.d[0] = inputs[0].d[0];
                out_dim.d[1] = exprBuilder.constant(1);
            }
            // detection_boxes
            else if (outputIndex == 1)
            {
                out_dim.nbDims = 3;
                out_dim.d[0] = inputs[0].d[0];
                out_dim.d[1] = numOutputBoxes;
                out_dim.d[2] = exprBuilder.constant(4);
            }
            // detection_scores: outputIndex == 2
            // detection_classes: outputIndex == 3
            else if (outputIndex == 2 || outputIndex == 3)
            {
                out_dim.nbDims = 2;
                out_dim.d[0] = inputs[0].d[0];
                out_dim.d[1] = numOutputBoxes;