TensorRT/parsers/caffe/caffeWeightFactory/caffeWeightFactory.h，caffeWeightFactory.cpp源碼研讀一

前言

TensorRT/parsers/caffe/caffeWeightFactory/caffeWeightFactory.h及TensorRT/parsers/caffe/caffeWeightFactory/caffeWeightFactory.cpp這兩個檔案裡定義了CaffeWeightFactory這個類別。

它被用來獲取、做型別轉換或隨機生成 神經網路的權重（nvinfer1::Weights）。此外還定義了一些metadata來記錄包括：是否初始化（mInitialized），轉換或獲取權重時是否出現異常（mOK），權重的資料型別（mDataType）等資訊。

由於篇幅限制，本篇先介紹建構子，圍繞著私有成員變數的函數，及checkForNans和sizeOfCaffeType這兩個utility function。

之後將繼續介紹用於轉換權重型別的函數，隨機生成權重的函數，與trtcaffe::BlobProto有關的函數及用於獲取權重的函數。

TensorRT/parsers/caffe/caffeWeightFactory/caffeWeightFactory.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_WEIGHT_FACTORY_H
#define TRT_CAFFE_PARSER_CAFFE_WEIGHT_FACTORY_H

#include <vector>
#include <string>
#include <random>
#include <memory>
#include "NvInfer.h"
#include "weightType.h"
#include "trtcaffe.pb.h"

namespace nvcaffeparser1
{
//用於生成nvinfer1::Weights的工廠？
//此類別用於獲取、做型別轉換或隨機生成 神經網路的權重（nvinfer1::Weights）
class CaffeWeightFactory
{
public:
    CaffeWeightFactory(const trtcaffe::NetParameter& msg, nvinfer1::DataType dataType, std::vector<void*>& tmpAllocs, bool isInitialized);
    nvinfer1::DataType getDataType() const;
    size_t getDataTypeSize() const;
    std::vector<void*>& getTmpAllocs();
    int getBlobsSize(const std::string& layerName);
    const trtcaffe::BlobProto* getBlob(const std::string& layerName, int index);
    std::vector<nvinfer1::Weights> getAllWeights(const std::string& layerName);
    virtual nvinfer1::Weights operator()(const std::string& layerName, WeightType weightType);
    void convert(nvinfer1::Weights& weights, nvinfer1::DataType targetType);
    void convert(nvinfer1::Weights& weights);
    bool isOK();
    bool isInitialized();
    /*
    class Weights
    定義於include/NvInferRuntime.h
    class Weights
	{
	public:
	    DataType type;      //!< The type of the weights.
	    const void* values; //!< The weight values, in a contiguous array.
	    int64_t count;      //!< The number of weights in the array.
	};
	用於表示神經網路內各層的權重
    */
    nvinfer1::Weights getNullWeights();
    nvinfer1::Weights allocateWeights(int64_t elems, std::uniform_real_distribution<float> distribution = std::uniform_real_distribution<float>(-0.01f, 0.01F));
    nvinfer1::Weights allocateWeights(int64_t elems, std::normal_distribution<float> distribution);
    static trtcaffe::Type getBlobProtoDataType(const trtcaffe::BlobProto& blobMsg);
    static size_t sizeOfCaffeType(trtcaffe::Type type);
    // The size returned here is the number of array entries, not bytes
    static std::pair<const void*, size_t> getBlobProtoData(const  trtcaffe::BlobProto& blobMsg, trtcaffe::Type type, std::vector<void*>& tmpAllocs);

private:
    template <typename T>
    bool checkForNans(const void* values, int count, const std::string& layerName);
    nvinfer1::Weights getWeights(const trtcaffe::BlobProto& blobMsg, const std::string& layerName);

    //記錄著神經網路相關的資訊，如：多少層，每層是什麼等
    const trtcaffe::NetParameter& mMsg;
    std::unique_ptr<trtcaffe::NetParameter> mRef;
    //將轉換權重或分配權重時申請的空間記錄於此
    std::vector<void*>& mTmpAllocs;
    //權重的資料型別
    nvinfer1::DataType mDataType;
    // bool mQuantize;
    bool mInitialized;
    //用於隨機生成權重的generator
    std::default_random_engine generator;
    //指出轉換或獲取權重是否異常
    bool mOK{true};
};
} //namespace nvcaffeparser1
#endif //TRT_CAFFE_PARSER_CAFFE_WEIGHT_FACTORY_H

TensorRT/parsers/caffe/caffeWeightFactory/caffeWeightFactory.cpp

/*
 * 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.
 */

#include "caffeMacros.h"
#include "caffeWeightFactory.h"
#include "half.h"

using namespace nvinfer1;
using namespace nvcaffeparser1;


//...


CaffeWeightFactory::CaffeWeightFactory(const trtcaffe::NetParameter& msg, DataType dataType, std::vector<void*>& tmpAllocs, bool isInitialized)
    : mMsg(msg)
    , mTmpAllocs(tmpAllocs)
    , mDataType(dataType)
    , mInitialized(isInitialized)
{
    mRef = std::unique_ptr<trtcaffe::NetParameter>(new trtcaffe::NetParameter);
}

//獲取初始化時就設定好的成員變數mDataType
DataType CaffeWeightFactory::getDataType() const
{
    return mDataType;
}

//獲取其數據類型所佔用的byte數
size_t CaffeWeightFactory::getDataTypeSize() const
{
    switch (getDataType())
    {
    case DataType::kFLOAT:
    case DataType::kINT32:
        return 4;
    case DataType::kHALF:
        return 2;
    case DataType::kINT8:
        return 1;
    }
    return 0;
}

//獲取初始化時就設定好的成員變數mTmpAllocs
std::vector<void*>& CaffeWeightFactory::getTmpAllocs()
{
    return mTmpAllocs;
}

bool CaffeWeightFactory::isOK()
{
    return mOK;
}

bool CaffeWeightFactory::isInitialized()
{
    return mInitialized;
}

/*
檢查values裡的東西能否有效地被轉成float（即轉成float後是否為NaN）
如果成功則回傳true，反之則回傳false
*/
template <typename T>
bool CaffeWeightFactory::checkForNans(const void* values, int count, const std::string& layerName)
{
    //先轉為T型別
    const T* v = reinterpret_cast<const T*>(values);
    for (int i = 0; i < count; i++)
    {
        if (std::isnan(float(v[i])))
        {
            std::cout << layerName << ": Nan detected in weights" << std::endl;
            return false;
        }
    }
    return true;
}

/*
trtcaffe::Type
定義於TensorRT/parsers/caffe/proto/trtcaffe.proto
// Math and storage types
enum Type {
  DOUBLE = 0;
  FLOAT = 1;
  FLOAT16 = 2;
  INT = 3;  // math not supported
  UINT = 4;  // math not supported
}
*/
//回傳caffe資料型別佔用的byte數
size_t CaffeWeightFactory::sizeOfCaffeType(trtcaffe::Type type)
{
    //分單精度，半精度，雙精度三種情況
    if (type == trtcaffe::FLOAT)
    {
        return sizeof(float);
    }
    if (type == trtcaffe::FLOAT16)
    {
        //uint16_t即unsigned short int，佔2 bytes
        return sizeof(uint16_t);
    }
    return sizeof(double);
}

std::isnan

在CaffeWeightFactory::checkForNans函數中用到了std::isnan函數，關於std::isnan，詳見C++ std::isnan及std::isinf。

uint16_t及int64_t

在sizeOfCaffeType 函數中用到了 uint16_t 這個型別，而 allocateWeights 函數則用到了 int64_t 這個型別，關於這些帶_t的型別，詳見C uint8_t，int64_t。

參考連結

C++ std::isnan及std::isinf

C uint8_t，int64_t