TensorRT/parsers/caffe/caffeParser/caffeParser.h源碼研讀

前言

caffeParser.h中宣告了CaffeParser這個類別。其具體實作則位於 caffeParser.cpp。

因為caffeParser.cpp內容較多，它的介紹共分為：
TensorRT/parsers/caffe/caffeParser/caffeParser.cpp入口函數源碼研讀， TensorRT/parsers/caffe/caffeParser/caffeParser.cpp - parse源碼研讀， TensorRT/parsers/caffe/caffeParser/caffeParser.cpp - parseBinaryProto源碼研讀， TensorRT/parsers/caffe/caffeParser/caffeParser.cpp - parseNormalizeParam源碼研讀 等幾篇。

TensorRT/parsers/caffe/caffeParser/caffeParser.h

/*
 * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef TRT_CAFFE_PARSER_CAFFE_PARSER_H
#define TRT_CAFFE_PARSER_CAFFE_PARSER_H

#include <vector>
#include <memory>
#include <unordered_map>
#include <string>

#include "NvCaffeParser.h"
#include "caffeWeightFactory.h"
#include "blobNameToTensor.h"
#include "trtcaffe.pb.h"

namespace nvcaffeparser1
{
//繼承自TensorRT/include/NvCaffeParser.h裡的抽象類別ICaffeParser
//ICafferParser裡宣告的都是virtual function，在這個ICafferParser的子類別中會賦予它們具體的定義
class CaffeParser : public ICaffeParser
{
public:
    const IBlobNameToTensor* parse(const char* deploy,
                                   const char* model,
                                   nvinfer1::INetworkDefinition& network,
                                   nvinfer1::DataType weightType) override;

    const IBlobNameToTensor* parseBuffers(const char* deployBuffer,
                                          size_t deployLength,
                                          const char* modelBuffer,
                                          size_t modelLength,
                                          nvinfer1::INetworkDefinition& network,
                                          nvinfer1::DataType weightType) override;

    void setProtobufBufferSize(size_t size) override { mProtobufBufferSize = size; }
    void setPluginFactory(nvcaffeparser1::IPluginFactory* factory) override { mPluginFactory = factory; }
    void setPluginFactoryExt(nvcaffeparser1::IPluginFactoryExt* factory) override
    {
        mPluginFactory = factory;
        mPluginFactoryIsExt = true;
    }

    void setPluginFactoryV2(nvcaffeparser1::IPluginFactoryV2* factory) override { mPluginFactoryV2 = factory; }
    void setPluginNamespace(const char* libNamespace) override { mPluginNamespace = libNamespace; }
    IBinaryProtoBlob* parseBinaryProto(const char* fileName) override;
    void destroy() override { delete this; }
    //以下二函數override ICaffeParser中的函數，但未實做
    //(void)recorder的作用是？
    void setErrorRecorder(nvinfer1::IErrorRecorder* recorder) override { (void)recorder; assert(!"TRT- Not implemented."); }
    nvinfer1::IErrorRecorder* getErrorRecorder() const override { assert(!"TRT- Not implemented."); return nullptr; }

private:
    ~CaffeParser() override;
    /*
    以下五個函數的作用都是將trtcaffe::LayerParameter裡的資料
    整理成std::vector<nvinfer1::PluginField>的格式後回傳
    */
    std::vector<nvinfer1::PluginField> parseNormalizeParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
    std::vector<nvinfer1::PluginField> parsePriorBoxParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
    std::vector<nvinfer1::PluginField> parseDetectionOutputParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
    std::vector<nvinfer1::PluginField> parseLReLUParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
    std::vector<nvinfer1::PluginField> parseRPROIParam(const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors);
    //分配size（默認為1）大小的T型別的空間，然後放到mTmpAllocs裡集中管理
    template <typename T>
    T* allocMemory(int size = 1)
    {
        T* tmpMem = static_cast<T*>(malloc(sizeof(T) * size));
        mTmpAllocs.push_back(tmpMem);
        return tmpMem;
    }

    const IBlobNameToTensor* parse(nvinfer1::INetworkDefinition& network,
                                   nvinfer1::DataType weightType,
                                   bool hasModel);

private:
    //用於表示模型架構
    std::shared_ptr<trtcaffe::NetParameter> mDeploy;
    //用於儲存模型權重
    std::shared_ptr<trtcaffe::NetParameter> mModel;
    /*
    為CaffeWeightFactory所用，
    用於記錄分配或轉換權重時，使用malloc新申請的記憶體
    */
    std::vector<void*> mTmpAllocs;
    /*
    BlobNameToTensor
    定義於TensorRT/parsers/caffe/blobNameToTensor.h
    其核心為mMap這個字典，用於把blob name對應到ITensor*
    */
    BlobNameToTensor* mBlobNameToTensor{nullptr};
    /*
    為TensorRT/parsers/caffe/caffeParser/readProto.h中的函數readBinaryProto所用，
    表示stream讀取的最大byte數
    */
    size_t mProtobufBufferSize{INT_MAX};
    /*
    nvcaffeparser1::IPluginFactory
    定義於TensorRT/include/NvCaffeParser.h
    Plugin factory used to configure plugins.
    */
    nvcaffeparser1::IPluginFactory* mPluginFactory{nullptr};
    /*
    nvcaffeparser1::IPluginFactoryV2
    定義於TensorRT/include/NvCaffeParser.h
    Plugin factory used to configure plugins.    
    */
    nvcaffeparser1::IPluginFactoryV2* mPluginFactoryV2{nullptr};
    bool mPluginFactoryIsExt{false};
    //用於存放nvinfer1::IPluginV2的向量
    std::vector<nvinfer1::IPluginV2*> mNewPlugins;
    //用於儲存的plugin名稱對應到nvinfer1::IPluginCreator*字典
    std::unordered_map<std::string, nvinfer1::IPluginCreator*> mPluginRegistry;
    //作為nvcaffeparser1::IPluginFactoryV2::createPlugin函數的參數使用
    std::string mPluginNamespace = "";
};
} //namespace nvcaffeparser1
#endif //TRT_CAFFE_PARSER_CAFFE_PARSER_H

delete this

destroy函數的內容是delete this，詳見C++ delete this。

std::unordered_map

在caffeParser.h中宣告了CaffeParser類別，其中有一個成員mPluginRegistry為std::unordered_map型別的：

std::unordered_map<std::string, nvinfer1::IPluginCreator*> mPluginRegistry;

std::map內部的元素會依其key被排序； 而std::unordered_map則沒有。如果要隨機存取一個元素，使用std::unordered_map的速度會比較快。

參考連結

TensorRT/parsers/caffe/caffeParser/caffeParser.cpp入口函數源碼研讀

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

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

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

C++ delete this

std::unordered_map