TensorRT/parsers/caffe/caffeParser/caffeParser.cpp - parseNormalizeParam源碼研讀

前言

TensorRT/parsers/caffe/caffeParser/caffeParser.cpp 中的 parse(INetworkDefinition&, DataType, bool) 函數用到了 parseNormalizeParam， parsePriorBoxParam， parseDetectionOutputParam， parseRPROIParam等函數，如：

std::vector<nvinfer1::PluginField> f;
if (layerMsg.type() == "Normalize")
{
    pluginName = "Normalize_TRT";
    f = parseNormalizeParam(layerMsg, weights, *mBlobNameToTensor);
}

本篇僅先就 parseNormalizeParam做介紹。

parseNormalizeParam

/*
將trtcaffe::LayerParameter::norm_param裡的資料整理成
std::vector<nvinfer1::PluginField>的格式後回傳
*/
std::vector<nvinfer1::PluginField> CaffeParser::parseNormalizeParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
	/*
	nvinfer1::PluginField
	定義於TensorRT/include/NvInferRuntimeCommon.h
	Structure containing plugin attribute field names and associated data
	This information can be parsed to decode necessary plugin metadata
	*/
    std::vector<nvinfer1::PluginField> f;
    /*
    trtcaffe::NormalizeParameter
    定義於TensorRT/parsers/caffe/proto/trtcaffe.proto
    Message that stores parameters used by NormalizeLayer
    */
    /*
    trtcaffe::LayerParameter::norm_param
    定義於TensorRT/parsers/caffe/proto/trtcaffe.proto
    */
    const trtcaffe::NormalizeParameter& p = msg.norm_param();

	/*
	allocMemory
	定義於TensorRT/parsers/caffe/caffeParser/caffeParser.h
	分配size（默認為1）大小的T型別的空間，然後放到mTmpAllocs裡集中管理
	template <typename T>
	T* allocMemory(int size = 1)
	*/
    int* acrossSpatial = allocMemory<int32_t>();
    *acrossSpatial = p.across_spatial() ? 1 : 0;
    /*
    PluginField的建構子：
    PluginField(const char* name_ = nullptr, const void* data_ = nullptr, const PluginFieldType type_ = PluginFieldType::kUNKNOWN, int32_t length_ = 0)
    */
    /*
    enum class PluginFieldType : int
	{
	    kFLOAT16 = 0, //!< FP16 field type.
	    kFLOAT32 = 1, //!< FP32 field type.
	    kFLOAT64 = 2, //!< FP64 field type.
	    kINT8 = 3,    //!< INT8 field type.
	    kINT16 = 4,   //!< INT16 field type.
	    kINT32 = 5,   //!< INT32 field type.
	    kCHAR = 6,    //!< char field type.
	    kDIMS = 7,    //!< nvinfer1::Dims field type.
	    kUNKNOWN = 8
	};
    */
    f.emplace_back("acrossSpatial", acrossSpatial, PluginFieldType::kINT32, 1);

    int* channelShared = allocMemory<int32_t>();
    *channelShared = p.channel_shared() ? 1 : 0;
    f.emplace_back("channelShared", channelShared, PluginFieldType::kINT32, 1);

    auto* eps = allocMemory<float>();
    *eps = p.eps();
    f.emplace_back("eps", eps, PluginFieldType::kFLOAT32, 1);

	/*
	Weights
	定義於TensorRT/include/NvInferRuntime.h
	An array of weights used as a layer parameter
	*/
    std::vector<Weights> w;
    // If .caffemodel is not provided, need to randomize the weight
    //如果權重尚未初始化，則生成正態分布的隨機權重
    if (!weightFactory.isInitialized())
    {
	    /*
	    parserutils::getCHW
	    定義於TensorRT/parsers/common/parserUtils.h
	    從Dims中抽取CHW三維的長度後回傳
	    */
        int C = parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions()).c();
        w.emplace_back(weightFactory.allocateWeights(C, std::normal_distribution<float>(0.0F, 1.0F)));
    }
    else
    {
        // Use the provided weight from .caffemodel
        w = weightFactory.getAllWeights(msg.name());
    }

    for (auto weight : w)
    {
        f.emplace_back("weights", weight.values, PluginFieldType::kFLOAT32, weight.count);
    }

    int* nbWeights = allocMemory<int32_t>();
    *nbWeights = w.size();
    f.emplace_back("nbWeights", nbWeights, PluginFieldType::kINT32, 1);

    return f;
}

parsePriorBoxParam

std::vector<nvinfer1::PluginField> CaffeParser::parsePriorBoxParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& /*tensors*/)
{
    std::vector<nvinfer1::PluginField> f;
    const trtcaffe::PriorBoxParameter& p = msg.prior_box_param();

    int minSizeSize = p.min_size_size();
    auto* minSize = allocMemory<float>(minSizeSize);
    for (int i = 0; i < minSizeSize; ++i)
    {
        minSize[i] = p.min_size(i);
    }
    f.emplace_back("minSize", minSize, PluginFieldType::kFLOAT32, minSizeSize);

    int maxSizeSize = p.max_size_size();
    auto* maxSize = allocMemory<float>(maxSizeSize);
    for (int i = 0; i < maxSizeSize; ++i)
    {
        maxSize[i] = p.max_size(i);
    }
    f.emplace_back("maxSize", maxSize, PluginFieldType::kFLOAT32, maxSizeSize);

    int aspectRatiosSize = p.aspect_ratio_size();
    auto* aspectRatios = allocMemory<float>(aspectRatiosSize);
    for (int i = 0; i < aspectRatiosSize; ++i)
    {
        aspectRatios[i] = p.aspect_ratio(i);
    }
    f.emplace_back("aspectRatios", aspectRatios, PluginFieldType::kFLOAT32, aspectRatiosSize);

    int varianceSize = p.variance_size();
    auto* variance = allocMemory<float>(varianceSize);
    for (int i = 0; i < varianceSize; ++i)
    {
        variance[i] = p.variance(i);
    }
    f.emplace_back("variance", variance, PluginFieldType::kFLOAT32, varianceSize);

    int* flip = allocMemory<int32_t>();
    *flip = p.flip() ? 1 : 0;
    f.emplace_back("flip", flip, PluginFieldType::kINT32, 1);

    int* clip = allocMemory<int32_t>();
    *clip = p.clip() ? 1 : 0;
    f.emplace_back("clip", clip, PluginFieldType::kINT32, 1);

    int* imgH = allocMemory<int32_t>();
    *imgH = p.has_img_h() ? p.img_h() : p.img_size();
    f.emplace_back("imgH", imgH, PluginFieldType::kINT32, 1);

    int* imgW = allocMemory<int32_t>();
    *imgW = p.has_img_w() ? p.img_w() : p.img_size();
    f.emplace_back("imgW", imgW, PluginFieldType::kINT32, 1);

    auto* stepH = allocMemory<float>();
    *stepH = p.has_step_h() ? p.step_h() : p.step();
    f.emplace_back("stepH", stepH, PluginFieldType::kFLOAT32, 1);

    auto* stepW = allocMemory<float>();
    *stepW = p.has_step_w() ? p.step_w() : p.step();
    f.emplace_back("stepW", stepW, PluginFieldType::kFLOAT32, 1);

    auto* offset = allocMemory<float>();
    *offset = p.offset();
    f.emplace_back("offset", offset, PluginFieldType::kFLOAT32, 1);

    return f;
}

parseDetectionOutputParam

std::vector<nvinfer1::PluginField> CaffeParser::parseDetectionOutputParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& /*tensors*/)
{
    std::vector<nvinfer1::PluginField> f;
    const trtcaffe::DetectionOutputParameter& p = msg.detection_output_param();
    const trtcaffe::NonMaximumSuppressionParameter& nmsp = p.nms_param();

    int* shareLocation = allocMemory<int32_t>();
    *shareLocation = p.share_location() ? 1 : 0;
    f.emplace_back("shareLocation", shareLocation, PluginFieldType::kINT32, 1);

    int* varianceEncodedInTarget = allocMemory<int32_t>();
    *varianceEncodedInTarget = p.variance_encoded_in_target() ? 1 : 0;
    f.emplace_back("varianceEncodedInTarget", varianceEncodedInTarget, PluginFieldType::kINT32, 1);

    int* backgroundLabelId = allocMemory<int32_t>();
    *backgroundLabelId = p.background_label_id();
    f.emplace_back("backgroundLabelId", backgroundLabelId, PluginFieldType::kINT32, 1);

    int* numClasses = allocMemory<int32_t>();
    *numClasses = p.num_classes();
    f.emplace_back("numClasses", numClasses, PluginFieldType::kINT32, 1);

    //nms
    int* topK = allocMemory<int32_t>();
    *topK = nmsp.top_k();
    f.emplace_back("topK", topK, PluginFieldType::kINT32, 1);

    int* keepTopK = allocMemory<int32_t>();
    /*
    Number of total bboxes to be kept per image after nms step.
    -1 means keeping all bboxes after nms step.
    */
    *keepTopK = p.keep_top_k();
    f.emplace_back("keepTopK", keepTopK, PluginFieldType::kINT32, 1);

    auto* confidenceThreshold = allocMemory<float>();
    *confidenceThreshold = p.confidence_threshold();
    f.emplace_back("confidenceThreshold", confidenceThreshold, PluginFieldType::kFLOAT32, 1);

    //nms
    auto* nmsThreshold = allocMemory<float>();
    *nmsThreshold = nmsp.nms_threshold();
    f.emplace_back("nmsThreshold", nmsThreshold, PluginFieldType::kFLOAT32, 1);

    // input order = {0, 1, 2} in Caffe
    int* inputOrder = allocMemory<int32_t>(3);
    inputOrder[0] = 0;
    inputOrder[1] = 1;
    inputOrder[2] = 2;
    f.emplace_back("inputOrder", inputOrder, PluginFieldType::kINT32, 3);

    // confSigmoid = false for Caffe
    int* confSigmoid = allocMemory<int32_t>();
    *confSigmoid = 0;
    f.emplace_back("confSigmoid", confSigmoid, PluginFieldType::kINT32, 1);

    // isNormalized = true for Caffe
    int* isNormalized = allocMemory<int32_t>();
    *isNormalized = 1;
    f.emplace_back("isNormalized", isNormalized, PluginFieldType::kINT32, 1);

    // codeTypeSSD : from NvInferPlugin.h
    // CORNER = 0, CENTER_SIZE = 1, CORNER_SIZE = 2, TF_CENTER = 3
    int* codeType = allocMemory<int32_t>();
    switch (p.code_type())
    {
    case trtcaffe::PriorBoxParameter::CORNER_SIZE:
        *codeType = static_cast<int>(plugin::CodeTypeSSD::CORNER_SIZE);
        break;
    case trtcaffe::PriorBoxParameter::CENTER_SIZE:
        *codeType = static_cast<int>(plugin::CodeTypeSSD::CENTER_SIZE);
        break;
    case trtcaffe::PriorBoxParameter::CORNER: // CORNER is default
    default:
        *codeType = static_cast<int>(plugin::CodeTypeSSD::CORNER);
        break;
    }
    f.emplace_back("codeType", codeType, PluginFieldType::kINT32, 1);

    return f;
}

parseLReLUParamparseRPROIParam

std::vector<nvinfer1::PluginField> CaffeParser::parseLReLUParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& /*tensors*/)
{
    std::vector<nvinfer1::PluginField> f;
    const trtcaffe::ReLUParameter& p = msg.relu_param();
    auto* negSlope = allocMemory<float>();
    *negSlope = p.negative_slope();
    f.emplace_back("negSlope", negSlope, PluginFieldType::kFLOAT32, 1);
    return f;
}

parseRPROIParam

std::vector<nvinfer1::PluginField> CaffeParser::parseRPROIParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
    std::vector<nvinfer1::PluginField> f;
    const trtcaffe::ROIPoolingParameter& p1 = msg.roi_pooling_param();
    const trtcaffe::RegionProposalParameter& p2 = msg.region_proposal_param();

    // Memory allocations for plugin field variables
    int* poolingH = allocMemory<int32_t>();
    int* poolingW = allocMemory<int32_t>();
    auto* spatialScale = allocMemory<float>();
    int* preNmsTop = allocMemory<int32_t>();
    int* nmsMaxOut = allocMemory<int32_t>();
    auto* iouThreshold = allocMemory<float>();
    auto* minBoxSize = allocMemory<float>();
    int* featureStride = allocMemory<int32_t>();
    int* anchorsRatioCount = allocMemory<int32_t>();
    int* anchorsScaleCount = allocMemory<int32_t>();
    int anchorsRatiosSize = p2.anchor_ratio_size();
    auto* anchorsRatios = allocMemory<float>(anchorsRatiosSize);
    int anchorsScalesSize = p2.anchor_scale_size();
    auto* anchorsScales = allocMemory<float>(anchorsScalesSize);

    // Intialize the plugin fields with values from the prototxt
    *poolingH = p1.pooled_h();
    f.emplace_back("poolingH", poolingH, PluginFieldType::kINT32, 1);

    *poolingW = p1.pooled_w();
    f.emplace_back("poolingW", poolingW, PluginFieldType::kINT32, 1);

    *spatialScale = p1.spatial_scale();
    f.emplace_back("spatialScale", spatialScale, PluginFieldType::kFLOAT32, 1);

    *preNmsTop = p2.prenms_top();
    f.emplace_back("preNmsTop", preNmsTop, PluginFieldType::kINT32, 1);

    *nmsMaxOut = p2.nms_max_out();
    f.emplace_back("nmsMaxOut", nmsMaxOut, PluginFieldType::kINT32, 1);

    *iouThreshold = p2.iou_threshold();
    f.emplace_back("iouThreshold", iouThreshold, PluginFieldType::kFLOAT32, 1);

    *minBoxSize = p2.min_box_size();
    f.emplace_back("minBoxSize", minBoxSize, PluginFieldType::kFLOAT32, 1);

    *featureStride = p2.feature_stride();
    f.emplace_back("featureStride", featureStride, PluginFieldType::kINT32, 1);

    *anchorsRatioCount = p2.anchor_ratio_count();
    f.emplace_back("anchorsRatioCount", anchorsRatioCount, PluginFieldType::kINT32, 1);

    *anchorsScaleCount = p2.anchor_scale_count();
    f.emplace_back("anchorsScaleCount", anchorsScaleCount, PluginFieldType::kINT32, 1);

    for (int i = 0; i < anchorsRatiosSize; ++i) {
        anchorsRatios[i] = p2.anchor_ratio(i);
}
    f.emplace_back("anchorsRatios", anchorsRatios, PluginFieldType::kFLOAT32, anchorsRatiosSize);

    for (int i = 0; i < anchorsScalesSize; ++i) {
        anchorsScales[i] = p2.anchor_scale(i);
}
    f.emplace_back("anchorsScales", anchorsScales, PluginFieldType::kFLOAT32, anchorsScalesSize);

    return f;
}