【玩转yolov5】使用TensorRT C++ API搭建yolov5s-v4.0网络结构(1)_yolov5s c++

注意：对于yolov5s-v4.0网络结构，上图仅作参考，实际结构以代码为准，存在少量差异!

#需要一个全局的ILogger对象,用于记录日志信息
static Logger gLogger; 
#创建一个网络生成器
IBuilder* builder = createInferBuilder(gLogger);
#使用IBuilder类方法创建一个空的网络
INetworkDefinition* network = builder->createNetworkV2(0U);

        builder是构建器，他会自动搜索cuda内核目录以获得最快的可用实现，构建和运行时的GPU需要保持一致。由builder构建的引擎(engine)不能跨平台和TensorRT版本移植。上面由builder创建了一个空的网络结构，后面就需要通过tensorrt c++ api来逐填充该网络结构，直至完整构建yolov5s-v4.0网络。

         首先构造focus结构，在yolov3和yolov4中并没有这个结构，其中比较关键的是切片操作。以我训练的输入为640*640*3的yolov5s的结构为例，原始640*640*3的图像输入focus结构，采用切片操作，生成320*320*12的特征图，再经过一个输出通道为32的卷积操作，生成320*320*32的特征图。focus结构的意义在于可以最大程度的减少信息损失而进行下采样操作。focus结构中需要用到的一个重要的tensorrt api就是addSlice接口，它用于创建一个slice层。

virtual ISliceLayer* nvinfer1::INetworkDefinition::addSlice(    ITenso&     input,
                                                                Dims     start,
                                                                Dims     size,
                                                                Dims     stride 
) 

https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_network_definition.html#ab3d64683b10afdbc944075da818fb086

ILayer* focus(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int inch, int outch, int ksize, std::string lname) {
    ISliceLayer *s1 = network->addSlice(input, Dims3{ 0, 0, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
    ISliceLayer *s2 = network->addSlice(input, Dims3{ 0, 1, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
    ISliceLayer *s3 = network->addSlice(input, Dims3{ 0, 0, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
    ISliceLayer *s4 = network->addSlice(input, Dims3{ 0, 1, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 });
    ITensor* inputTensors[] = { s1->getOutput(0), s2->getOutput(0), s3->getOutput(0), s4->getOutput(0) };
    auto cat = network->addConcatenation(inputTensors, 4); #通道维度上的拼接
    auto conv = convBlock(network, weightMap, *cat->getOutput(0), outch, ksize, 1, 1, lname + ".conv");
    return conv;
}

接下来是一个CBL结构，这个比较好理解，拆开来看就是：Conv + BN + Silu。注意，虽然上面的全局网络结构图中展示的CBL中的激活函数是LeakyRelu，但是在v4.0中激活函数是Silu(Sigmoid Weighted Linear Unit)，是一种较为平滑的激活函数。

ILayer* convBlock(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname) {
    Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
    int p = ksize / 2;
    IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], emptywts);
    assert(conv1);
    conv1->setStrideNd(DimsHW{ s, s });
    conv1->setPaddingNd(DimsHW{ p, p });
    conv1->setNbGroups(g);
    IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-3);
 
    // silu = x * sigmoid
    auto sig = network->addActivation(*bn1->getOutput(0), ActivationType::kSIGMOID);
    assert(sig);
    auto ew = network->addElementWise(*bn1->getOutput(0), *sig->getOutput(0), ElementWiseOperation::kPROD);
    assert(ew);
    return ew;
}

因为后面要频繁用到该结构，这里拆开来详细讲解一下。首先是卷积，调用addConvolutionNd来创建一个新的卷积层。 因为没有bias一项，定义的bias的Weights结构中values为nullptr。stride,padding,group等参数通过IConvolutionLayer的内部成员函数来设置。

    Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
    int p = ksize / 2;
    IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], emptywts);
    assert(conv1);
    conv1->setStrideNd(DimsHW{ s, s });
    conv1->setPaddingNd(DimsHW{ p, p });
    conv1->setNbGroups(g);

然后是BN层，回顾一下BN层的定义：

E [ x ] 是batch的均值，V a r [ x ] 是batch的方差，ϵ为了防止除0，γ 对应batch学习得到的权重，β 就是偏置。

TensorRT中并没有直接的BatchNorm层，该层实际上是通过转换系数依靠Scale层来完成。

好了，万事具备，可以手撕代码了。

IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) {
    float *gamma = (float*)weightMap[lname + ".weight"].values;
    float *beta = (float*)weightMap[lname + ".bias"].values;
    float *mean = (float*)weightMap[lname + ".running_mean"].values; //均值
    float *var = (float*)weightMap[lname + ".running_var"].values; //方差
    int len = weightMap[lname + ".running_var"].count;
 
    //scale
    float *scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
    for (int i = 0; i < len; i++) {
        scval[i] = gamma[i] / sqrt(var[i] + eps);
    }
    Weights scale{ DataType::kFLOAT, scval, len };
 
    //shift
    float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
    for (int i = 0; i < len; i++) {
        shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
    }
    Weights shift{ DataType::kFLOAT, shval, len };
 
    //power
    float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
    for (int i = 0; i < len; i++) {
        pval[i] = 1.0;
    }
    Weights power{ DataType::kFLOAT, pval, len };
 
    weightMap[lname + ".scale"] = scale;
    weightMap[lname + ".shift"] = shift;
    weightMap[lname + ".power"] = power;
    //BatchNorm是channel维度的操作
    IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
    assert(scale_1);
    return scale_1;
}

然后就是激活函数Silu，从下面的公式可以看出来其实就是给sigmoid激活函数加了一个权重，这个权重恰恰就是输入。

f(x)=x⋅σ(x)    

f′(x)=f(x)+σ(x)(1−f(x))

同样，TensorRT中也没有直接提供Silu的api，通过addActivation配合addElementWise中的乘操作可以轻松构建Silu。

 // silu = x * sigmoid
    auto sig = network->addActivation(*bn1->getOutput(0), ActivationType::kSIGMOID);
    assert(sig);
    auto ew = network->addElementWise(*bn1->getOutput(0), *sig->getOutput(0), ElementWiseOperation::kPROD);
    assert(ew);

【参考文献】

https://zhuanlan.zhihu.com/p/172121380