TensorRT/parsers/caffe/caffeParser/opParsers/parseConv.cpp源碼研讀

TensorRT/parsers/caffe/caffeParser/opParsers/parseConv.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 "opParsers.h"

using namespace nvinfer1;

namespace nvcaffeparser1
{
/*
獲取或隨機初始化權重及偏置量後，把它們記錄於mTmpAllocs中
然後為網路新增卷積層
*/
ILayer* parseConvolution(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& weightFactory, BlobNameToTensor& tensors)
{
    //卷積層的輸入及輸出個數均應為1
    if (!checkBlobs(msg, 1, 1))
    {
        return nullptr;
    }

    /*
    trtcaffe::ConvolutionParameter
    定義於TensorRT/parsers/caffe/proto/trtcaffe.proto
    */
    const trtcaffe::ConvolutionParameter& p = msg.convolution_param();
    int nbOutputs = p.num_output();

    int kernelH = p.has_kernel_h() ? p.kernel_h() : p.kernel_size(0);
    //如果未指定kernel_w，則默認它與kernel_h一致
    int kernelW = p.has_kernel_w() ? p.kernel_w() : p.kernel_size_size() > 1 ? p.kernel_size(1) : p.kernel_size(0);
    int C = parserutils::getCHW(tensors[msg.bottom(0)]->getDimensions()).c();
    // The group size for group conv
    int G = p.has_group() ? p.group() : 1;

    auto CbyG = float(C / G * nbOutputs);
    float std_dev = 1.0F / sqrtf((kernelW * kernelH * sqrtf(CbyG)));
    /*
    WeightType::kGENERIC，WeightType::kBIAS：
    types for convolution, deconv, fully connected
    分別代表權重及偏置量
    */
    /*
    CaffeWeightFactory::operator()
    定義於TensorRT/parsers/caffe/caffeWeightFactory/caffeWeightFactory.cpp
    獲取layerName層的第int(weightType)個blob（即權重或偏置量），
    如果沒有，則回傳長度為0的Null weights
    */
    /*
    CaffeWeightFactory::allocateWeights
    Weights CaffeWeightFactory::allocateWeights(int64_t elems, std::normal_distribution<float> distribution)
    定義於TensorRT/parsers/caffe/caffeWeightFactory/caffeWeightFactory.cpp
    分配getDataType()型別，elems大小的空間，用它來新建Weights物件後回傳
    採用normal distribution來隨機初始化
    */
    /*
    CaffeWeightFactory::getNullWeights
    定義於TensorRT/parsers/caffe/caffeWeightFactory/caffeWeightFactory.cpp
    回傳空的weight array
    */
    //獲取或隨機初始化權重及偏置量
    //權重張量（四維）的體積為kernelW * kernelH * CbyG
    Weights kernelWeights = weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kGENERIC) : weightFactory.allocateWeights(kernelW * kernelH * CbyG, std::normal_distribution<float>(0.0F, std_dev));
    Weights biasWeights = !p.has_bias_term() || p.bias_term() ? (weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kBIAS) : weightFactory.allocateWeights(nbOutputs)) : weightFactory.getNullWeights();

    /*
    CaffeWeightFactory::convert
    將weights裡的內容轉換為初始化時就設定好的mDataType型別，
    並將weights.values記錄於mTmpAllocs中
    */
    weightFactory.convert(kernelWeights);
    weightFactory.convert(biasWeights);
    auto inTensor = tensors[msg.bottom(0)];
    /*
    INetworkDefinition::addConvolution
    宣告於TensorRT/include/NvInfer.h
    TRT_DEPRECATED virtual IConvolutionLayer* addConvolution(ITensor& input, int nbOutputMaps, DimsHW kernelSize,
        Weights kernelWeights, Weights biasWeights) TRTNOEXCEPT = 0;
    */
    /*
    nvinfer1::DimsHW
    定義於TensorRT/include/NvInfer.h
    */
    auto layer = network.addConvolution(*inTensor, nbOutputs, DimsHW{kernelH, kernelW}, kernelWeights, biasWeights);

    if (layer)
    {
        //stride：默認為1
        int strideH = p.has_stride_h() ? p.stride_h() : p.stride_size() > 0 ? p.stride(0) : 1;
        int strideW = p.has_stride_w() ? p.stride_w() : p.stride_size() > 1 ? p.stride(1) : p.stride_size() > 0 ? p.stride(0) : 1;

        //pad：默認為0
        int padH = p.has_pad_h() ? p.pad_h() : p.pad_size() > 0 ? p.pad(0) : 0;
        int padW = p.has_pad_w() ? p.pad_w() : p.pad_size() > 1 ? p.pad(1) : p.pad_size() > 0 ? p.pad(0) : 0;

        //dilation：默認為1
        int dilationH = p.dilation_size() > 0 ? p.dilation(0) : 1;
        int dilationW = p.dilation_size() > 1 ? p.dilation(1) : p.dilation_size() > 0 ? p.dilation(0) : 1;

        /*
        宣告於TensorRT/include/NvInfer.h
        TRT_DEPRECATED virtual void setStride(DimsHW stride) TRTNOEXCEPT = 0;
        TRT_DEPRECATED virtual void setPadding(DimsHW padding) TRTNOEXCEPT = 0;
        virtual void setPaddingMode(PaddingMode paddingMode) TRTNOEXCEPT = 0;
        TRT_DEPRECATED virtual void setDilation(DimsHW dilation) TRTNOEXCEPT = 0;
        */
        layer->setStride(DimsHW{strideH, strideW});
        layer->setPadding(DimsHW{padH, padW});
        /*
        定義於TensorRT/include/NvInfer.h
        enum class PaddingMode : int
		{
		    kEXPLICIT_ROUND_DOWN = 0, //!< Use explicit padding, rounding output size down.
		    kEXPLICIT_ROUND_UP = 1,   //!< Use explicit padding, rounding output size up.
		    kSAME_UPPER = 2,          //!< Use SAME padding, with prePadding <= postPadding.
		    kSAME_LOWER = 3,          //!< Use SAME padding, with prePadding >= postPadding.
		    kCAFFE_ROUND_DOWN = 4,    //!< Use CAFFE padding, rounding output size down, uses prePadding value.
		    kCAFFE_ROUND_UP = 5       //!< Use CAFFE padding, rounding output size up, uses prePadding value.
		};
		*/
		layer->setPaddingMode(PaddingMode::kCAFFE_ROUND_DOWN);
        layer->setDilation(DimsHW{dilationH, dilationW});

		/*
		宣告於TensorRT/include/NvInfer.h
		virtual void setNbGroups(int nbGroups) TRTNOEXCEPT = 0;
		*/
        layer->setNbGroups(G);
    }
    return layer;
}
} //namespace nvcaffeparser1

std::normal_distribution

在函數parseConvolution中用到了std::normal_distribution：

/**/std::normal_distribution<float>(0.0F, std_dev)/**/

關於std::normal_distribution，詳見C++ uniform_real_distribution及normal_distribution。

weight initialization

auto CbyG = float(C / G * nbOutputs);
float std_dev = 1.0F / sqrtf((kernelW * kernelH * sqrtf(CbyG)));
Weights kernelWeights = weightFactory.isInitialized() ? weightFactory(msg.name(), WeightType::kGENERIC) : weightFactory.allocateWeights(kernelW * kernelH * CbyG, std::normal_distribution<float>(0.0F, std_dev));