五. TensorRT API的基本使用-build-trt-module

前言

自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考

本次课程我们来学习课程第五章—TensorRT API 的基本使用，一起来学习 TensorRT 模型搭建的模块化

课程大纲可以看下面的思维导图

0. 简述

本小节目标：学习 TensorRT 模型搭建的模块化设计思想

今天我们来讲第五章节第六小节—5.6-build-sub-graph 这个案例，上节课中我们学习了利用 C++ API 搭建网络的各个 layer 比如 conv、bn、reshape 等等，其实我们就可以利用这些 layer 搭建一个完整的模型，例如 resnet、yolo、vit 等等，理论上我们都是可以实现的

但是整个模型的搭建其实还是有非常多的地方需要注意的，我们目前只学习了搭建某个 layer，那这些都是模型的一些小零件，如果按照一个个 layer 去搭建其实是非常麻烦的，因为一个模型可能就有上百层，一层层搭建太过于繁琐，整个代码也非常冗长，所以我们这里学习 pytorch 搭建的方式采用模块化设计的思想，先搭建一个个 module 然后利用这些 module 搭建 network

下面我们开始本次课程的学习🤗

1. 案例运行

在正式开始课程之前，博主先带大家跑通 5.6-build-sub-graph 这个小节的案例🤗
源代码获取地址：https://github.com/kalfazed/tensorrt_starter

首先大家需要把 tensorrt_starter 这个项目给 clone 下来，指令如下：

git clone https://github.com/kalfazed/tensorrt_starter.git

也可手动点击下载，点击右上角的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击 here 下载博主准备好的源代码（注意代码下载于 2024/7/14 日，若有改动请参考最新）

整个项目后续需要使用的软件主要有 CUDA、cuDNN、TensorRT、OpenCV，大家可以参考 Ubuntu20.04软件安装大全 进行相应软件的安装，博主这里不再赘述

假设你的项目、环境准备完成，下面我们来一起运行 5.6 小节案例代码

开始之前我们需要创建几个文件夹，在 tensorrt_starter/chapter5-tensorrt-api-basics/5.6-build-sub-graph 小节中韩君老师已经创建了 models 文件夹，并且在 models 文件夹下有提供的 weights 文件夹，我们只需要在 models 文件夹下新建一个 onnx 和 engine 文件夹即可

创建完后 5.6 小节整个目录结构如下：

虽然这里已经提供了各种 weights 文件，但是我们还是执行下对应的 python 脚本，方便对比 python 和 c++ 的结果，先进入到 5.6 小节中：

cd tensorrt_starter/chapter5-tensorrt-api-basics/5.6-build-sub-graph

执行如下指令：

python src/python/export_cbr.py

Note：大家需要准备一个虚拟环境，安装好 torch、onnx、onnxsim 等第三方库

输出如下：

生成好的 onnx 模型文件保存在 models/onnx 文件夹下，对应的 weights 文件保存在 models/weights 文件夹下，大家可以查看

接着我们需要加载 weights 利用自己搭建的 network 生成对应的 engine，在此之前我们需要修改下整体的 Makefile.config，指定一些库的路径：

# tensorrt_starter/config/Makefile.config
# CUDA_VER                    :=  11
CUDA_VER                    :=  11.6
    
# opencv和TensorRT的安装目录
OPENCV_INSTALL_DIR          :=  /usr/local/include/opencv4
# TENSORRT_INSTALL_DIR        :=  /mnt/packages/TensorRT-8.4.1.5
TENSORRT_INSTALL_DIR        :=  /home/jarvis/lean/TensorRT-8.6.1.6

Note：大家查看自己的 CUDA 是多少版本，修改为对应版本即可，另外 OpenCV 和 TensorRT 修改为你自己安装的路径即可

然后我们还要简单修改下源码，在 src/cpp/main.cpp 中默认使用的 weights 是 sample_c2f.weights，我们修改为 sample_cbr.weights，修改如下所示：

# src/cpp/main.cpp
int main(int argc, char const *argv[])
{
    Model model("models/weights/sample_cbr.weights", Model::precision::FP32);
    // Model model("models/weights/sample_c2f.weights", Model::precision::FP16);
    ...
}

接着我们就可以来执行编译，指令如下：

make -j64

输出如下：

接着执行：

./trt-infer

输出如下：

我们这里通过手动构建一个 conv+bn+relu 的 module 并加载相应的 weights 权重完成模块的构建和推理，可以看到和 python 推理结果保持一致

Note：博主这里也准备了其它的模块和相应的权重，大家可以点击 here 下载，然后运行代码看下其它模块的搭建过程

如果大家能够看到上述输出结果，那就说明本小节案例已经跑通，下面我们就来看看具体的代码实现

2. 代码分析

2.1 main.cpp

我们先从 main.cpp 看起：

#include <iostream>
#include <memory>

#include "utils.hpp"
#include "model.hpp"

using namespace std;

int main(int argc, char const *argv[])
{
    /*
     * 这里面依次举几个例子来进行展示, 对应的输入和输出也会不一样
     * sample_cbr:            conv + BN + ReLU:                input shape: [1x1x5x5],     output shape: [1x1x3x3]
     * sample_resBlock:       ---:                             input shape: [1x1x5x5],     output shape: [1x3x5x5]
     * sample_convBNSiLU:     conv + BN + SeLU:                input shape: [1x1x5x5],     output shape: [1x3x5x5]
     * sample_c2f:            ---:                             input shape: [1x1x5x5],     output shape: [1x4x5x5]
    */
    
    // Model model("models/weights/sample_cbr.weights", Model::precision::FP32);
    // Model model("models/weights/sample_cbr.weights", Model::precision::FP16);
    // Model model("models/weights/sample_resBlock.weights", Model::precision::FP32);
    // Model model("models/weights/sample_resBlock.weights", Model::precision::FP16);
    // Model model("models/weights/sample_convBNSiLU.weights", Model::precision::FP32);
    // Model model("models/weights/sample_c2f.weights", Model::precision::FP32);
    Model model("models/weights/sample_c2f.weights", Model::precision::FP16);

    if(!model.build()){
        LOGE("fail in building model");
        return 0;
    }
    if(!model.infer()){
        LOGE("fail in infering model");
        return 0;
    }
    return 0;
}

Model 的构造函数和上节课差不多，接收一个 weights 此外还额外添加了一个 precision 精度参数，可以设置为 FP32 或者 FP16，我们简单看下其构造函数的实现：

Model::Model(string path, precision prec){
    if (getFileType(path) == ".onnx")
        mOnnxPath = path;
    else if (getFileType(path) == ".weights")
        mWtsPath = path;
    else 
        LOGE("ERROR: %s, wrong weight or model type selected. Program terminated", getFileType(path).c_str());

    if (prec == precision::FP16) {
        mPrecision = nvinfer1::DataType::kHALF;
    } else if (prec == precision::INT8) {
        mPrecision = nvinfer1::DataType::kINT8;
    } else {
        mPrecision = nvinfer1::DataType::kFLOAT;
    }

    mEnginePath = getEnginePath(path, prec);
}

其中 precision 精度是一个枚举类型，根据传入的 prec 参数来指定 mPrecision 精度，其类型是 nvinfer1::DataType

2.2 model.cpp

我们重点来看下 build 接口：

bool Model::build() {
    if (mOnnxPath != "") {
        return build_from_onnx();
    } else {
        return build_from_weights();
    }
}

和上小节案例一样，也是调用的 build_from_weights 函数来 build model

我们重点来看下该函数的实现：

if (fileExists(mEnginePath)){
    LOG("%s has been generated!", mEnginePath.c_str());
    return true;
} else {
    LOG("%s not found. Building engine...", mEnginePath.c_str());
}

mWts = loadWeights();

// 这里和之前的创建方式是一样的
Logger logger;
auto builder       = make_unique<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
auto config        = make_unique<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
auto network       = make_unique<nvinfer1::INetworkDefinition>(builder->createNetworkV2(1));

前面的这部分和上小节案例一模一样，都是先加载 weights 然后创建 build，创建 config，创建 network

// 根据不同的网络架构创建不同的TensorRT网络，这里使用几个简单的例子
if (mWtsPath == "models/weights/sample_cbr.weights") {
    network::build_cbr(*network, mPrecision, mWts);
} else if (mWtsPath == "models/weights/sample_resBlock.weights") {
    network::build_resBlock(*network, mPrecision, mWts);
} else if (mWtsPath == "models/weights/sample_convBNSiLU.weights") {
    network::build_convBNSiLU(*network, mPrecision, mWts);
} else if (mWtsPath == "models/weights/sample_c2f.weights") {
    network::build_C2F(*network, mPrecision, mWts);
} else {
    return false;
}

接着根据不同的 module 调用不同的函数，只不过这里通过 network 命名空间来创建的，下面我们来看看 network 这个命名空间所做的工作

2.3 network.hpp

network.hpp 头文件代码如下：

#ifndef __NETWORK_HPP__
#define __NETWORK_HPP__

#include <NvInfer.h>
#include <string>
#include <map>
#include <memory>
#include <model.hpp>

namespace network {


namespace parser {

nvinfer1::IShuffleLayer* addReshape(
    std::string layer_name,
    nvinfer1::ITensor& input, 
    std::vector<int> dims,
    std::vector<int> perm,
    nvinfer1::INetworkDefinition& network);

nvinfer1::IShuffleLayer* addPermute(
    std::string layer_name,
    nvinfer1::ITensor& input, 
    std::vector<int> perm,
    nvinfer1::INetworkDefinition& network);

nvinfer1::IFullyConnectedLayer* addFullyConnected(
    std::string layer_name,
    nvinfer1::ITensor& input, 
    int output_channel,
    nvinfer1::INetworkDefinition& network,
    std::map<std::string, nvinfer1::Weights> weights);

nvinfer1::IScaleLayer* addBatchNorm(
    std::string layer_name,
    nvinfer1::ITensor& input,
    nvinfer1::INetworkDefinition& network,
    std::map<std::string, nvinfer1::Weights> weights);

nvinfer1::IConvolutionLayer* addConv2d(
    std::string layer_name, 
    nvinfer1::ITensor& input,
    int kernel_size, int output_channel, int stride, int pad,
    nvinfer1::DataType prec,
    nvinfer1::INetworkDefinition& network,
    std::map<std::string, nvinfer1::Weights> weights);

nvinfer1::IActivationLayer* addActivation(
    std::string layer_name,
    nvinfer1::ITensor& input, 
    nvinfer1::ActivationType type,
    nvinfer1::INetworkDefinition& network);

nvinfer1::IElementWiseLayer* addElementWise(
    std::string layer_name,
    nvinfer1::ITensor& input1, 
    nvinfer1::ITensor& input2, 
    nvinfer1::ElementWiseOperation type,
    nvinfer1::INetworkDefinition& network);

nvinfer1::IElementWiseLayer* addConvBNSiLU(
    std::string layer_name, 
    nvinfer1::ITensor& input, 
    int kernel_size, 
    int output_channel, 
    int stride, 
    int pad,
    nvinfer1::DataType prec,
    nvinfer1::INetworkDefinition& network,
    std::map<std::string, nvinfer1::Weights> weights);

nvinfer1::ILayer* addC2F(
    std::string layer_name, 
    nvinfer1::ITensor& input, 
    int output_channel, 
    nvinfer1::DataType prec,
    nvinfer1::INetworkDefinition& network,
    std::map<std::string, nvinfer1::Weights> weights);


} // namespace parser

// void build_linear(nvinfer1::INetworkDefinition& network, std::map<std::string, nvinfer1::Weights> mWts);
// void build_conv(nvinfer1::INetworkDefinition& network, std::map<std::string, nvinfer1::Weights> mWts);
// void build_permute(nvinfer1::INetworkDefinition& network, std::map<std::string, nvinfer1::Weights> mWts);
// void build_reshape(nvinfer1::INetworkDefinition& network, std::map<std::string, nvinfer1::Weights> mWts);
// void build_batchNorm(nvinfer1::INetworkDefinition& network, std::map<std::string, nvinfer1::Weights> mWts);

void build_cbr(
    nvinfer1::INetworkDefinition& network, 
    nvinfer1::DataType prec,
    std::map<std::string, nvinfer1::Weights> weights) ;

void build_resBlock(
    nvinfer1::INetworkDefinition& network,
    nvinfer1::DataType prec,
    std::map<std::string, nvinfer1::Weights> weights) ;

void build_convBNSiLU(
    nvinfer1::INetworkDefinition& network,
    nvinfer1::DataType prec,
    std::map<std::string, nvinfer1::Weights> weights) ;

void build_C2F(
    nvinfer1::INetworkDefinition& network,
    nvinfer1::DataType prec,
    std::map<std::string, nvinfer1::Weights> weights);


}; // namespace network

#endif //__NETWORK_HPP__

这个头文件定义了 network 命名空间，并包含了一系列函数声明，这些函数用来在 TensorRT 中构建不同的模块，此外这个命名空间下还定义了一个子命名空间 parser，用于在 TensorRT 网络定义中添加各种 layer 和 module

下面我们就来看不同模块案例的具体实现

3. 案例

3.1 sample_cbr

我们先看 sample_cbr 案例，代码如下：

void build_cbr(
    nvinfer1::INetworkDefinition& network, 
    nvinfer1::DataType prec,
    map<string, nvinfer1::Weights> weights) 
{
    auto input  = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});

    auto conv   = parser::addConv2d("conv", *input, 3, 3, 1, 0, prec, network, weights);
    auto bn     = parser::addBatchNorm("norm", *conv->getOutput(0), network, weights);
    auto leaky  = parser::addActivation("leaky", *bn->getOutput(0), nvinfer1::ActivationType::kLEAKY_RELU, network);

    leaky->getOutput(0) ->setName("output0");
    network.markOutput(*leaky->getOutput(0));
}

首先通过 network 的 addInput 接口创建输入向量，接着调用 parser 命名空间下的模块函数 addConv2d、addBatchNorm、addActivation 分别添加 conv、bn 和 leaky relu，与上个案例相比，conv、bn 和 relu 的具体实现被封装到了 parser 下面，使得代码更加简洁和模块化

下面是 parser 中各个 layer 的具体实现：

nvinfer1::IScaleLayer* addBatchNorm(
    string layer_name,
    nvinfer1::ITensor& input,
    nvinfer1::INetworkDefinition& network,
    std::map<std::string, nvinfer1::Weights> weights)
{
    // 因为TensorRT内部没有BatchNorm的实现，但是我们只要知道BatchNorm的计算原理，就可以使用IScaleLayer来创建BN的计算
    // IScaleLayer主要是用在quantization和dequantization，作为提前了解，我们试着使用IScaleLayer来搭建于一个BN的parser
    // IScaleLayer可以实现: y = (x * scale + shift) ^ pow

    float* gamma   = (float*)weights[layer_name + ".weight"].values;
    float* beta    = (float*)weights[layer_name + ".bias"].values;
    float* mean    = (float*)weights[layer_name + ".running_mean"].values;
    float* var     = (float*)weights[layer_name + ".running_var"].values;
    float  eps     = 1e-5;
    
    int    count   = weights[layer_name + ".running_var"].count;

    float* scales  = (float*)malloc(count * sizeof(float));
    float* shifts  = (float*)malloc(count * sizeof(float));
    float* pows    = (float*)malloc(count * sizeof(float));
    
    // 这里具体参考一下batch normalization的计算公式，网上有很多
    for (int i = 0; i < count; i ++) {
        scales[i] = gamma[i] / sqrt(var[i] + eps);
        shifts[i] = beta[i] - (mean[i] * gamma[i] / sqrt(var[i] + eps));
        pows[i]   = 1.0;
    }

    // 将计算得到的这些值写入到Weight中
    auto scales_weights = nvinfer1::Weights{nvinfer1::DataType::kFLOAT, scales, count};
    auto shifts_weights = nvinfer1::Weights{nvinfer1::DataType::kFLOAT, shifts, count};
    auto pows_weights   = nvinfer1::Weights{nvinfer1::DataType::kFLOAT, pows, count};

    // 创建IScaleLayer并将这些weights传进去，这里使用channel作为scale model
    auto bn = network.addScale(input, nvinfer1::ScaleMode::kCHANNEL, shifts_weights, scales_weights, pows_weights);
    bn->setName(layer_name.c_str());

    LOGV("%s, %s", bn->getName(), (printDims(bn->getOutput(0)->getDimensions())).c_str());

    return bn;
}

nvinfer1::IConvolutionLayer* addConv2d(
    string layer_name, 
    nvinfer1::ITensor& input, 
    int kernel_size, 
    int output_channel, 
    int stride, 
    int pad,
    nvinfer1::DataType prec,
    nvinfer1::INetworkDefinition& network,
    std::map<std::string, nvinfer1::Weights> weights)
{
    auto conv = network.addConvolutionNd(
            input, output_channel, 
            nvinfer1::DimsHW{kernel_size, kernel_size}, 
            weights[layer_name + ".weight"], 
            weights[layer_name + ".bias"]);
    conv->setName(layer_name.c_str());
    conv->setStride(nvinfer1::DimsHW(stride, stride));
    conv->setPaddingNd(nvinfer1::DimsHW(pad, pad));

    // 注意，这里setPrecision需要跟config->setFlag配合使用，否则无效
    conv->setPrecision(prec);
    LOGV("%s, %s", conv->getName(), (printDims(conv->getOutput(0)->getDimensions())).c_str());

    return conv;
}

nvinfer1::IActivationLayer* addActivation(
    string layer_name,
    nvinfer1::ITensor& input, 
    nvinfer1::ActivationType type,
    nvinfer1::INetworkDefinition& network)
{
    auto act = network.addActivation(input, type);
    act->setName(layer_name.c_str());
    LOGV("%s, %s", act->getName(), (printDims(act->getOutput(0)->getDimensions())).c_str());

    return act;
}

这些函数封装了 TensorRT 底层 API 的调用，其实也就是上节课的各个 layer 的案例，这种方式简化了模型构建过程，存在以下优点：

• 模块化设计：通过将各个网络层的添加逻辑封装成函数，提高了代码的可读性和可维护性
• 易于使用：用户只需调用这些函数并提供必要的参数，即可轻松得向 TensorRT 网络定义中添加相应的层
• 扩展性：这种设计模式允许轻松地扩展其他类型的网络层或功能，满足不同的模型需求

该案例执行后的输出如下所示：

FP16 精度的输出如下所示：

对比下 python 结果：

可以看到输出基本都相同，这个就是 sample_cbr 案例

3.2 sample_resBlock

下面我们来看看 sample_resBlock 案例，代码如下：

// 做一个residual block: 
//       conv0
//       /   \
//      /    conv1
//     |      |
//     |      bn1
//     |      |
//     |     relu1
//     |      |
//     |      |
//     |     conv2
//     |      |
//     |      bn2
//      \    /
//       \  /
//       add2
//        |
//       relu2
//
void build_resBlock(
    nvinfer1::INetworkDefinition& network,
    nvinfer1::DataType prec,
    map<string, nvinfer1::Weights> weights) 
{
    auto data  = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});

    auto conv0 = parser::addConv2d("conv0", *data, 3, 3, 1, 1, prec, network, weights);

    auto conv1 = parser::addConv2d("conv1", *conv0->getOutput(0), 3, 3, 1, 1, prec, network, weights);
    auto bn1   = parser::addBatchNorm("norm1", *conv1->getOutput(0), network, weights);
    auto relu1 = parser::addActivation("relu1", *bn1->getOutput(0), nvinfer1::ActivationType::kRELU, network);

    auto conv2 = parser::addConv2d("conv2", *relu1->getOutput(0), 3, 3, 1, 1, prec, network, weights);
    auto bn2   = parser::addBatchNorm("norm2", *conv2->getOutput(0), network, weights);

    auto add2  = parser::addElementWise("add2", *conv0->getOutput(0), *bn2->getOutput(0), nvinfer1::ElementWiseOperation::kSUM, network);
    auto relu2 = parser::addActivation("relu2", *add2->getOutput(0), nvinfer1::ActivationType::kRELU, network);

    relu2->getOutput(0) ->setName("output0");
    network.markOutput(*relu2->getOutput(0));
}

build_resBlock 函数用于在 TensorRT 中构建一个 Residual Block 残差块，我们知道其结构后就可以像搭积木一样把它组合起来就行了，其中 conv、bn、activation 我们上个案例已经分析过了，这里我们分析下 ElementWise 这个 layer，具体实现如下：

nvinfer1::IElementWiseLayer* addElementWise(
    string layer_name,
    nvinfer1::ITensor& input1, 
    nvinfer1::ITensor& input2,
    nvinfer1::ElementWiseOperation type,
    nvinfer1::INetworkDefinition& network)
{
    auto ew = network.addElementWise(input1, input2, type);
    ew->setName(layer_name.c_str());
    LOGV("%s, %s", ew->getName(), (printDims(ew->getOutput(0)->getDimensions())).c_str());
    
    return ew;
}

addElementWise 函数用于在 TensorRT 中添加一个 IElementWiseLayer 层，该层执行元素级操作（element-wise operations），即在逐元素基础上对输入张量进行操作。这种操作对于实现诸如加法、乘法、最小值、最大值等功能非常有用，尤其是在实现残差连接时，其中：

• layer_name：用于标识该层的名称
• input1: 第一个输入张量
• input2: 第二个输入张量
• type: 元素级操作的类型，由 nvinfer1::ElementWiseOperation 枚举类型指定，它包括
  • kSUM: 逐元素加法
  • kPROD: 逐元素乘法
  • …

该案例执行后的输出如下所示：

FP16 精度的输出如下所示：

对比下 python 结果：

可以看到输出基本都相同，这个就是 sample_resBlock 案例

3.3 sample_convBNSiLU

下面我们来看看 sample_convBNSiLU 案例，代码如下：

// 做一个conv + bn + SiLU: (yolov8的模块测试)
//        conv
//         |
//         bn
//       /   \
//      |     |
//      |    sigmoid
//      \     /
//       \   /
//        Mul
//

void build_convBNSiLU(
    nvinfer1::INetworkDefinition& network,
    nvinfer1::DataType prec,
    map<string, nvinfer1::Weights> weights) 
{
    auto data  = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});

    auto silu  = parser::addConvBNSiLU("", *data, 3, 3, 1, 1, prec, network, weights);


    silu->getOutput(0) ->setName("output0");
    network.markOutput(*silu->getOutput(0));
}

build_convBNSiLU 函数用于在 TensorRT 中构建一个包含卷积（Conv）、批量归一化（BatchNorm）和 SiLU 激活（Swish 激活函数）的模块，这是 YOLOv8 网络结构中的一个常见模块。它主要是通过 parser 中的 addConvBNSiLU 函数来实现的，代码如下：

nvinfer1::IElementWiseLayer* addConvBNSiLU(
    string layer_name, 
    nvinfer1::ITensor& input, 
    int kernel_size, 
    int output_channel, 
    int stride, 
    int pad,
    nvinfer1::DataType prec,
    nvinfer1::INetworkDefinition& network,
    std::map<std::string, nvinfer1::Weights> weights)
{
    auto conv    = addConv2d(layer_name + "conv", input, kernel_size, output_channel, stride, pad, prec, network, weights);
    auto bn      = addBatchNorm(layer_name + "norm", *conv->getOutput(0), network, weights);
    auto sigmoid = addActivation(layer_name + "sigmoid", *bn->getOutput(0), nvinfer1::ActivationType::kSIGMOID, network);
    auto mul     = addElementWise(layer_name + "mul", *bn->getOutput(0), *sigmoid->getOutput(0), nvinfer1::ElementWiseOperation::kPROD, network);

    return mul;
} 

实现也非常简单，其中的 conv、bn、activation 以及 elementwise 我们前面都已经讲过了，值得注意的是这里的 ElementWiseOperation 类型 kPROD 也就是我们前面提到的逐元素相乘，这是因为 SiLU（SWish） 激活函数的公式是 $SiLU(x)=x\cdot \sigma (x)$，其中的 $\sigma (x)$ 是 Sigmodi 函数

该案例执行后的输出如下所示：

对比下 python 结果：

可以看到输出基本都相同，这个就是 sample_convBNSiLU 案例

3.4 sample_c2f

下面我们来看看 sample_c2f 案例，代码如下：

// 做一个C2F: (yolov8的模块测试)
//        input
//          |
//      convBNSiLU (n * ch)
//       /  |  \
//      /   |   \
//     |    |    |
//     |    | convBNSiLU ( 0.5n * ch)
//     |    |    |
//     |    | convBNSiLU ( 0.5n * ch)
//     |    |    |
//     |    \    /
//     |     \  /
//     |     add (0.5n * ch)
//      \    /
//       \  /
//      Concat (1.5n * ch)
//        |
//    convBNSiLU (n * ch)

void build_C2F(
    nvinfer1::INetworkDefinition& network,
    nvinfer1::DataType prec,
    map<string, nvinfer1::Weights> weights) 
{
    auto data  = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});

    auto c2f  = parser::addC2F("", *data, 4, prec, network, weights);

    c2f->getOutput(0) ->setName("output0");
    network.markOutput(*c2f->getOutput(0));
}

C2f 也是 YOLOv8 网络结构中的一个常见模块，它主要是通过 parser 中的 addC2F 函数来实现的，它的结构如下图所示：

具体实现代码如下：

nvinfer1::ILayer* addC2F(
    string layer_name, 
    nvinfer1::ITensor& input, 
    int output_channel, 
    nvinfer1::DataType prec,
    nvinfer1::INetworkDefinition& network,
    std::map<std::string, nvinfer1::Weights> weights)
{
    auto cv1     = addConvBNSiLU(layer_name + "cv1.", input, 1, output_channel, 1, 0, prec, network, weights);
    auto dim     = cv1->getOutput(0)->getDimensions();

    auto slice1  = addSlice(layer_name + "slice1", 
                            *cv1->getOutput(0), 
                            nvinfer1::Dims4{0,        0,          0,        0},         // B, C, H, W (0, 0, 0, 0)
                            nvinfer1::Dims4{dim.d[0], dim.d[1]/2, dim.d[2], dim.d[3]},  // B, 1/2 * C, H, W
                            nvinfer1::Dims4{1,        1,          1,        1},         // 1, 1, 1, 1
                            network);
    auto slice2  = addSlice(layer_name + "slice2", 
                            *cv1->getOutput(0), 
                            nvinfer1::Dims4{0,        dim.d[1]/2, 0,        0},         // B, C, H, W (0, 1/2 * C, 0, 0)
                            nvinfer1::Dims4{dim.d[0], dim.d[1]/2, dim.d[2], dim.d[3]},  // B, 1/2 * C, H, W
                            nvinfer1::Dims4{1,        1,          1,        1},         // 1, 1, 1, 1
                            network);

    auto add     = addBottleNeck(layer_name + "m.0.", *slice2->getOutput(0), 2, 2, true, prec, network, weights);

    nvinfer1::ITensor* concat2Input[] = {cv1->getOutput(0), add->getOutput(0)};
    auto concat2 = addConcat(layer_name + "concat2", concat2Input, 2, network);

    auto cv2     = addConvBNSiLU(layer_name + "cv2.", *concat2->getOutput(0), 1, output_channel, 1, 0, prec, network, weights);

    return cv2;
}

C2F 模块中包含以下组件：

• ConvBNSiLU
• Slice
• BottleNeck
• Concat

其中的 ConvBNSiLU 我们已经分析过了，我们来看看其它三个组件，先看 slice 切片操作：

nvinfer1::ISliceLayer* addSlice(
    string layer_name,
    nvinfer1::ITensor& input, 
    nvinfer1::Dims start,
    nvinfer1::Dims size,
    nvinfer1::Dims stride,
    nvinfer1::INetworkDefinition& network)
{
    auto slice = network.addSlice(input, start, size, stride);
    slice->setName(layer_name.c_str());
    LOGV("%s, %s", slice->getName(), (printDims(slice->getOutput(0)->getDimensions())).c_str());

    return slice;
}

addSlice 函数的主要功能是将输入张量按照指定的参数进行切片操作，并返回创建的切片层，具体是通过 network.addSlice 方法完成，其参数解释如下：

• input：输入张量
• start：指定切片的起始位置
• size：指定切片后输出的维度
• stride：指定切片的步长

我们再回过头来看看 C2F 模块中的两个 slice 操作具体做了些什么，代码如下：

auto slice1  = addSlice(layer_name + "slice1", 
                        *cv1->getOutput(0), 
                        nvinfer1::Dims4{0,        0,          0,        0},         // B, C, H, W (0, 0, 0, 0)
                        nvinfer1::Dims4{dim.d[0], dim.d[1]/2, dim.d[2], dim.d[3]},  // B, 1/2 * C, H, W
                        nvinfer1::Dims4{1,        1,          1,        1},         // 1, 1, 1, 1
                        network);
auto slice2  = addSlice(layer_name + "slice2", 
                        *cv1->getOutput(0), 
                        nvinfer1::Dims4{0,        dim.d[1]/2, 0,        0},         // B, C, H, W (0, 1/2 * C, 0, 0)
                        nvinfer1::Dims4{dim.d[0], dim.d[1]/2, dim.d[2], dim.d[3]},  // B, 1/2 * C, H, W
                        nvinfer1::Dims4{1,        1,          1,        1},         // 1, 1, 1, 1
                        network);

其中的 slice1 我们后续没用到，我们重点来看 slice2 操作：

• 输入张量：*cv1->getOutput(0)，即卷积层 cv1 的输出张量。
• 起始位置 (start)：nvinfer1::Dims4{0, dim.d[1]/2, 0, 0}
  • 从批次维度的起始位置开始（0）
  • 从通道维度的中间位置开始（dim.d[1]/2），即从通道数的一半位置开始
  • 从高度维度的起始位置开始（0）
  • 从宽度维度的起始位置开始（0）
• 切片大小 (size)：nvinfer1::Dims4{dim.d[0], dim.d[1]/2, dim.d[2], dim.d[3]}
  • dim.d[0]: 保持批次维度的大小不变
  • dim.d[1]/2: 通道数的一半，即切出剩下的一半通道
  • dim.d[2]: 保持高度维度的大小不变
  • dim.d[3]: 保持宽度维度的大小不变
• 步长 (stride)：nvinfer1::Dims4{1, 1, 1, 1}
  • 在每个维度上按步长 1 进行切片操作

slice2 的输出回作为 BottleNeck 的输入，下面我们就一起来看下 BottleNeck 的实现：

// 做一个bottleneck: (yolov8的模块测试)
//    input
//    /  \
//   |    |
//   | convBNSiLU ( 0.5n * ch)
//   |    |
//   | convBNSiLU ( 0.5n * ch)
//   |    |
//   \    /
//    \  /
//    add (0.5n * ch)

nvinfer1::ILayer* addBottleNeck(
    string layer_name, 
    nvinfer1::ITensor& input, 
    int ch1, int ch2,
    bool shortcut,
    nvinfer1::DataType prec,
    nvinfer1::INetworkDefinition& network,
    std::map<std::string, nvinfer1::Weights> weights)
{
    auto silu1 = addConvBNSiLU(layer_name + "cv1.", input,                3, ch1, 1, 1, prec, network, weights);
    auto silu2 = addConvBNSiLU(layer_name + "cv2.", *silu1->getOutput(0), 3, ch2, 1, 1, prec, network, weights);

    if (shortcut)  {
        auto add  =  addElementWise(layer_name + "cv1.add", 
                                    input, *silu2->getOutput(0),  
                                    nvinfer1::ElementWiseOperation::kSUM, network);
        return add;
    }
    
    return silu1;
}

该函数包含以下组件：

• ConvBNSiLU
• ElementWise
• 返回值则根据是否启用 shortcut 连接返回不同的输出层

我们来看下 C2F 模块中的 BottleNeck：

auto add     = addBottleNeck(layer_name + "m.0.", *slice2->getOutput(0), 2, 2, true, prec, network, weights);

其各个参数含义如下：

• layer_name + “m.0.”：层名称
• *slice2->getOutput(0)：slice2 的输出张量，作为 bottleneck 层的输入
• 2：第一个卷积层的输出通道数
• 2：第二个卷积层的输出通道数
• true：启用 shortcut 连接
• prec：精度
• network：网络定义对象
• weights：权重映射

我们再来看下 Concat 的实现：

nvinfer1::IConcatenationLayer* addConcat(
    string layer_name,
    nvinfer1::ITensor* input[],
    int size,
    nvinfer1::INetworkDefinition& network)
{
    auto concat =  network.addConcatenation(input, size);
    concat->setName(layer_name.c_str());
    LOGV("%s, %s", concat->getName(), (printDims(concat->getOutput(0)->getDimensions())).c_str());

    return concat;
}

该函数在 TensorRT 网络定义中添加一个拼接层，将多个输入张量按通道维度进行拼接，并返回创建的拼接层，具体实现是使用 network.addConcatenation 方法，其参数解释如下：

• input：输入张量组，包含需要拼接的多个张量
• size：输入张量的数量

最后我们再来看下 C2F 模块的 Concat：

nvinfer1::ITensor* concat2Input[] = {cv1->getOutput(0), add->getOutput(0)};
auto concat2 = addConcat(layer_name + "concat2", concat2Input, 2, network);

在 C2F 模块中，Concat 操作用于将 cv1 和 add 的输出张量拼接在一起，也就是第一个 ConvBNSiLU 模块的输出和 BottleNeck 的输出拼接在一起，最后再经过一个 convBNSiLU 就组成了我们的 C2F 模块

该案例执行后的输出如下所示：

FP16 精度的输出如下所示：

对比下 python 结果：

可以看到输出基本都相同，这个就是 sample_c2f 案例

实际上一个网络模型就是这些不同的模块组成的，我们后续在做 INT8 量化时就可以这样逐层分析，看具体哪一些层对精度影响比较严重，逐层或者说逐模块分析，最后把这些对精度影响比较严重的层不让它做量化，让它跑 FP16 的精度就行

总结

本次课程我们主要学习了 TensorRT 模型构建中模块化的思想，这里以 resBlock 以及 YOLOv8 中的几个常见模块为例去讲解了如何搭建一个模块，那通过这些模块我们就可以自己搭建整个网络框架。当然这种模块化的思想也可以用于 INT8 量化时精度下降的原因排查

OK，以上就是 5.6 小节案例的全部内容了，下节我们来学习 5.7 小节自定义插件的构建，敬请期待😄

下载链接

• tensorrt_starter源码
• 5.6-build-trt-module案例文件

参考

• Ubuntu20.04软件安装大全
• https://github.com/kalfazed/tensorrt_starter.git