五. TensorRT API的基本使用-build-model-from-scratch

前言

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

本次课程我们来学习课程第五章—TensorRT API 的基本使用，一起来学习利用 C++ API 手动搭建 network

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

0. 简述

本小节目标：学习利用 C++ API 从头开始搭建 network

今天我们来讲第五章节第五小节—5.5-build-model 这个案例，我们前面 build model 都是通过 onnxparser 解析器去 parse 我们导出好的 onnx 模型，这节我们来学习利用 C++ API 自己搭建一个 network 完成模型的 build 过程

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

1. 案例运行

在正式开始课程之前，博主先带大家跑通 5.5-build-model 这个小节的案例🤗
源代码获取地址：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.5 小节案例代码

开始之前我们需要创建几个文件夹，在 tensorrt_starter/chapter5-tensorrt-api-basics/5.5-build-model 小节中创建一个 models 文件夹，接着在 models 文件夹下创建 onnx 和 engine 和 weights 文件夹，总共四个文件夹需要创建

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

接着我们需要执行 python 文件创建一个 ONNX 模型并将其 weights 给保存下来，先进入到 5.5 小节中：

cd tensorrt_starter/chapter5-tensorrt-api-basics/5.5-build-model

执行如下指令：

python src/python/export_linear.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_sclice.weights，我们修改为 sample_linear.weights，修改如下所示：

# src/cpp/main.cpp
int main(int argc, char const *argv[])
{
    Model model("models/weights/sample_linear.weights");
    // Model model("models/weights/sample_slice.weights");
    ...
}

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

make -j64

输出如下：

接着执行：

./trt-infer

输出如下：

我们这里通过手动构建的 network 并加载相应的 weights 权重完成模型的构建和推理，可以看到和 python 推理结果保持一致，我们的模型使用的是一个简单的只包含 linear 层的 network

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_linear:         linear only:                     input shape: [1x5],         output shape: [1]
     * sample_conv:           conv only:                       input shape: [1x1x5x5],     output shape: [1x3x3x3]
     * sample_permute:        conv + permute:                  input shape: [1x1x5x5],     output shape: [1x3x3x3]
     * sample_reshape:        conv + reshape + linear:         input shape: [1x1x5x5],     output shape: [1x9x3]
     * sample_batchNorm:      conv + batchNorm:                input shape: [1x1x5x5],     output shape: [1x3x3x3]
     * sample_cbr:            conv + BN + ReLU:                input shape: [1x1x5x5],     output shape: [1x1x3x3]
    */

    // Model model("models/weights/sample_linear.weights");
    // Model model("models/weights/sample_conv.weights");
    // Model model("models/weights/sample_permute.weights");
    // Model model("models/weights/sample_reshape.weights");
    // Model model("models/weights/sample_batchNorm.weights");
    // Model model("models/weights/sample_cbr.weights");
    // Model model("models/weights/sample_pooling.weights");
    // Model model("models/weights/sample_upsample.weights");
    // Model model("models/weights/sample_deconv.weights");
    // Model model("models/weights/sample_concat.weights");
    // Model model("models/weights/sample_elementwise.weights");
    // Model model("models/weights/sample_reduce.weights");
    Model model("models/weights/sample_slice.weights");

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

与之前 build 的案例不同，我们这里传入到 model 中的是对应的 weights 权重，然后通过 model.build 接口构建 engine，通过 model.infer 接口完成推理

这里韩君老师提供了非常多的模型 build 的案例，大家感兴趣的可以多测试测试

2.2 model.cpp

我们重点来看下 build 接口发生了哪些变化：

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

我们可以看到如果 mOnnxPath 不为空则通过 build_from_onnx 函数来 build model，也就是我们之前案例所做的，如果 mOnnxPath 为空则通过 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();

首先我们通过 loadWeights 函数将对应的权重加载，分析该函数之前我们先看 python 是如何将 weights 给保存下来的，对应的代码如下：

def export_weight(model):
    current_path = os.path.dirname(__file__)
    f = open(current_path + "/../../models/weights/sample_linear.weights", 'w')
    f.write("{}\n".format(len(model.state_dict().keys())))
    
    # 我们将权重里的float数据，按照hex16进制的形式进行保存，也就是所谓的编码
    # 可以使用python中的struct.pack
    for k,v in model.state_dict().items():
        print('exporting ... {}: {}'.format(k, v.shape))
        
        # 将权重转为一维
        vr = v.reshape(-1).cpu().numpy()
        f.write("{} {}".format(k, len(vr)))
        for vv in vr:
            f.write(" ")
            f.write(struct.pack(">f", float(vv)).hex())
        f.write("\n")

为了能够让 TensorRT 读取 PyTorch 导出的权重，我们可以把权重按照指定的格式导出：

• count
• [name][len][weights value in hex mode]
• …

count 代表总的权重数量，之后的每一行代表一个 weight，最开始是 weight 的名字 name，接着是它的长度 len，接着是它的数据 value，注意这里的 value 是以 16 进制的格式保存下来的

我们会遍历整个 model 的所有参数，然后将参数 reshape 为一维数组，接着将参数的名称和一维数组的长度写入文件，最后将每个权重值转换为 float 格式，使用 struct.pack 将 float 数据转换为二进制数据，并将其转换为十六进制字符串格式保存

保存下来的 weights 类似于下面这种格式：

8
conv.weight 27 be578f59 3d5de7fd 3c4bcbc7 3d2a83cd be04920c bf03231a be586bd2 3f0d08f7 bc005dd1 3f243af1 be908d47 3d8a930a bef7665f bbe3706a be8e998f be2627c3 be8c1e94 bd6ac825 bb069f5f bef8ff71 3ee06550 be4430e3 bd4e884f be8ad2f9 be67ab35 be007c56 be1c17f5
conv.bias 3 3e3986f6 bea6d5da 3e53997a
norm.weight 3 3f866666 3f866666 3f866666
norm.bias 3 3d4ccccd 3d4ccccd 3d4ccccd
norm.running_mean 3 00000000 00000000 00000000
norm.running_var 3 3f800000 3f800000 3f800000
norm.num_batches_tracked 1 00000000
linear.weight 5 be3a4dc8 3f1e43aa 3d20cfdd 3f6d1bf0 bf75ab18

我们知道了 python 是怎么保存 weights 之后我们再来看下 c++ 是如何加载的，代码如下：

map<string, nvinfer1::Weights> Model::loadWeights(){
    ifstream f;
    if (!fileExists(mWtsPath)){ 
        LOGE("ERROR: %s not found", mWtsPath.c_str());
    }

    f.open(mWtsPath);

    int32_t size;
    map<string, nvinfer1::Weights> maps;
    f >> size;

    if (size <= 0) {
        LOGE("ERROR: no weights found in %s", mWtsPath.c_str());
    }

    while (size > 0) {
        nvinfer1::Weights weight;
        string name;
        int weight_length;

        f >> name;
        f >> std::dec >> weight_length;

        uint32_t* values = (uint32_t*)malloc(sizeof(uint32_t) * weight_length);
        for (int i = 0; i < weight_length; i ++) {
            f >> std::hex >> values[i];
        }

        weight.type = nvinfer1::DataType::kFLOAT;
        weight.count = weight_length;
        weight.values = values;

        maps[name] = weight;

        size --;
    }

    return maps;
}

我们需要将权重存储在一个 map<string, nvinfer1::Weights> 结构中，以便在构建神经网络时使用这些权重。

首先我们先检查下权重文件路径是否存在，接着从文件中读取权重的数量 size，然后循环读取权重数据，首先读取权重的名称 name 和权重的长度 weight_length，接着分配一个 uint32_t 的数组用于存储权重，数组大小为 weights_length

随后使用十六进制格式读取每个权重，并存储在 values 数据中，并将之前创建的 nvinfer1::Weights 变量进行一些设置：

• weight.type：数据类型
• weight.count：权重的数量
• weight.values：指向权重数据的指针

设置完成之后将权重存储在 mapes 中，键为权重名称，值为 nvinfer1::Weights 结构，最后返回包含所有权重的 maps

我们再回到 build 函数中：

Logger logger;
auto builder       = make_unique<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
auto network       = make_unique<nvinfer1::INetworkDefinition>(builder->createNetworkV2(1));
auto config        = make_unique<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());

load_weights 之后就和前面的 build 案例差不多，定义 Logger 创建 build，通过 builder 创建 network，创建 config

if (mWtsPath == "models/weights/sample_linear.weights") {
    build_linear(*network, mWts);
} else if (mWtsPath == "models/weights/sample_conv.weights") {
    build_conv(*network, mWts);
} else if (mWtsPath == "models/weights/sample_permute.weights") {
    build_permute(*network, mWts);
} else if (mWtsPath == "models/weights/sample_reshape.weights") {
    build_reshape(*network, mWts);
} else if (mWtsPath == "models/weights/sample_batchNorm.weights") {
    build_batchNorm(*network, mWts);
} else if (mWtsPath == "models/weights/sample_cbr.weights") {
    build_cbr(*network, mWts);
} else if (mWtsPath == "models/weights/sample_pooling.weights") {
    build_pooling(*network, mWts);
} else if (mWtsPath == "models/weights/sample_upsample.weights") {
    build_upsample(*network, mWts);
} else if (mWtsPath == "models/weights/sample_deconv.weights") {
    build_deconv(*network, mWts);
} else if (mWtsPath == "models/weights/sample_concat.weights") {
    build_concat(*network, mWts);
} else if (mWtsPath == "models/weights/sample_elementwise.weights") {
    build_elementwise(*network, mWts);
} else if (mWtsPath == "models/weights/sample_reduce.weights") {
    build_reduce(*network, mWts);
} else if (mWtsPath == "models/weights/sample_slice.weights") {
    build_slice(*network, mWts);
} else {
    return false;
}

这里的代码就与之前有所不同，我们会根据不同的网络架构加载不同的 weights 创建不同的 TensorRT 网络，之前我们是直接通过 onnxparser 进行 ONNX 模型解析的：

auto parser        = make_unique<nvonnxparser::IParser>(nvonnxparser::createParser(*network, logger));

if (!parser->parseFromFile(mOnnxPath.c_str(), 1)){
    LOGE("ERROR: failed to %s", mOnnxPath.c_str());
    return false;
}

我们来看下 build_linear 函数具体是怎么创建一个 network 的呢？

void Model::build_linear(nvinfer1::INetworkDefinition& network, map<string, nvinfer1::Weights> mWts) {
    auto data          = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 1, 5});
    auto fc            = network.addFullyConnected(*data, 1, mWts["linear.weight"], {});
    fc->setName("linear1");

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

整个网络结构如下：

/*
 * network 
 *
 *  -- input --    ITensor
 *  ---- | ----
 *  ---linear--    Ilayer
 *  ---- | ----
 *  -- output -    ITensor
*/

网络非常简单包含只包含一个 linear 层，它的类型是 Ilayer，它有一个输入和一个输出，类型是 ITensor，所以我们除了 linear 层外还需要创建输入和输出

首先我们通过 network.addInput 创建一个输入，其中：

• input0 是输入张量的名称
• nvinfer1::DataType::kFLOAT 指定输入数据的类型为浮点数
• nvinfer1::Dims4{1, 1, 1, 5} 指定输入张量的维度
  • Dims4 表示一个四维张量，这里的维度是 (1, 1, 1, 5)，通常表示 BxCxHxW

接着通过 network.addFullyConnected 创建一个 linear 层，其中：

• *data 是输入张量
• 1 是输出张量的通道数，即全连接层的输出大小，在这里，输出是一个单一的值
• mWts[“linear.weight”] 提供了该全连接层的权重
  • 这些权重是从之前加载的权重映射 mWts 中获取的，"linear.weight" 是权重的键，mWts["linear.weight"] 返回一个 nvinfer1::Weights 对象，包含全连接层的权重
• {} 是偏置项，空的偏置表示没有偏置项，或者偏置项为零

然后通过 setName 设置全连接层的名称，最后设置输出张量的名称并标记为网络输出：

• fc->getOutput(0)->setName(“output0”) 设置全连接层的输出张量的名称为 "output0"
• network.markOutput(*fc->getOutput(0)) 将这个输出张量标记为网络的输出，这意味着它是最终的输出，并且在推理时会输出这个张量的值。

在代码中我们可以看到通过 network.addXXX 可以添加某个 layer 层，那具体 TensorRT 支持哪些 layer 呢？其实我们可以查看它的官方文档，如下图所示：

更多细节大家可以查看：nvinfer1::ILayer Class Reference

那 network build 之后接下来的事情其实和前面的案例差不多：

config->setMaxWorkspaceSize(1<<28);
builder->setMaxBatchSize(1);

auto engine        = make_unique<nvinfer1::ICudaEngine>(builder->buildEngineWithConfig(*network, *config));
auto plan          = builder->buildSerializedNetwork(*network, *config);
auto runtime       = make_unique<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));

auto f = fopen(mEnginePath.c_str(), "wb");
fwrite(plan->data(), 1, plan->size(), f);
fclose(f);

mEngine            = shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(plan->data(), plan->size()), InferDeleter());
mInputDims         = network->getInput(0)->getDimensions();
mOutputDims        = network->getOutput(0)->getDimensions();

// 把优化前和优化后的各个层的信息打印出来
LOG("Before TensorRT optimization");
print_network(*network, false);
LOG("");
LOG("After TensorRT optimization");
print_network(*network, true);

// 最后把map给free掉
for (auto& mem : mWts) {
    free((void*) (mem.second.values));
}
LOG("Finished building engine");
return true;

通过 network 创建 engine，接着序列化，保存文件，最后把 map 给释放掉

以上就是手动构建 network 的过程，下面我们来看 infer 推理部分

我们在 infer 需要做的事情主要有：

• 1. 读取 model，创建 runtime，engine，context
• 2. 将数据从 host 传输到 device
• 3. 使用 context 推理
• 4. 将推理完的数据从 device 传输到 host

其实和前面的案例差不多，整体代码如下：

bool Model::infer(){

    /* 1. 读取model => 创建runtime, engine, context */
    if (!fileExists(mEnginePath)) {
        LOGE("ERROR: %s not found", mEnginePath.c_str());
        return false;
    }

    vector<unsigned char> modelData;
    modelData = loadFile(mEnginePath);
    
    Logger logger;
    auto runtime     = make_unique<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));
    auto engine      = make_unique<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(modelData.data(), modelData.size()));
    auto context     = make_unique<nvinfer1::IExecutionContext>(engine->createExecutionContext());

    auto input_dims   = context->getBindingDimensions(0);
    auto output_dims  = context->getBindingDimensions(1);

    LOG("input dim shape is:  %s", printDims(input_dims).c_str());
    LOG("output dim shape is: %s", printDims(output_dims).c_str());

    /* 2. 创建流 */
    cudaStream_t stream;
    cudaStreamCreate(&stream);

    /* 2. 初始化input，以及在host/device上分配空间 */
    init_data(input_dims, output_dims);

    /* 2. host->device的数据传递*/
    cudaMemcpyAsync(mInputDevice, mInputHost, mInputSize, cudaMemcpyKind::cudaMemcpyHostToDevice, stream);

    /* 3. 模型推理, 最后做同步处理 */
    float* bindings[] = {mInputDevice, mOutputDevice};
    bool success = context->enqueueV2((void**)bindings, stream, nullptr);

    /* 4. device->host的数据传递 */
    cudaMemcpyAsync(mOutputHost, mOutputDevice, mOutputSize, cudaMemcpyKind::cudaMemcpyDeviceToHost, stream);
    cudaStreamSynchronize(stream);

    LOG("input data is:  %s", printTensor(mInputHost, mInputSize / sizeof(float), input_dims).c_str());
    LOG("output data is: %s", printTensor(mOutputHost, mOutputSize / sizeof(float), output_dims).c_str());
    LOG("finished inference");
    return true;
}

那以上就是 sample_linear 案例的 build 和 infer 的完整过程了，下面我们再看看其它几个案例

3. 案例

3.1 sample_conv

我们来看下 sample_conv 案例，它通过 build_conv 函数来搭建 conv 网络，代码如下：

void Model::build_conv(nvinfer1::INetworkDefinition& network, map<string, nvinfer1::Weights> mWts) {
    auto data          = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});
    auto conv          = network.addConvolutionNd(*data, 3, nvinfer1::DimsHW{3, 3}, mWts["conv.weight"], mWts["conv.bias"]);
    conv->setName("conv1");
    conv->setStride(nvinfer1::DimsHW(1, 1));

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

和 sample_linear 案例一样，也是先 addInput 创建输入，接着通过 addConvolutionNd 创建 conv layer，然后设置 conv layer 的名称以及 stride，最后设置输出张量名称并标记为网络输出

其中 network.addConvolutionNd 的参数主要有：

• *data 是输入张量
• 3 是输出通道数
• nvinfer1::DimsHW{3, 3} 指定卷积核的大小
• mWts[“conv.weight”] 和 mWts[“conv.bias”] 分别提供卷积核的权重和偏置。

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

对比下 python 结果：

可以看到输出数据都相同，这个就是 sample_conv 案例

3.2 sample_permute

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

void Model::build_permute(nvinfer1::INetworkDefinition& network, map<string, nvinfer1::Weights> mWts) {
    auto data          = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});
    auto conv          = network.addConvolutionNd(*data, 3, nvinfer1::DimsHW{3, 3}, mWts["conv.weight"], mWts["conv.bias"]);
    conv->setName("conv1");
    conv->setStride(nvinfer1::DimsHW(1, 1));
    
    auto permute       = network.addShuffle(*conv->getOutput(0));
    permute->setFirstTranspose(nvinfer1::Permutation{0, 2, 3, 1}); // B, C, H, W -> B, H, W, C
    permute->setName("permute1");

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

permute 的实现我们主要是通过下面两行代码实现的：

auto permute       = network.addShuffle(*conv->getOutput(0));
permute->setFirstTranspose(nvinfer1::Permutation{0, 2, 3, 1}); // B, C, H, W -> B, H, W, C

其中：

• network.addShuffle 方法用于在网络中添加一个转置（permute）层。
• *conv->getOutput(0) 是卷积层的输出张量，作为转置层的输入。
• permute->setFirstTranspose(nvinfer1::Permutation{0, 2, 3, 1}) 设置转置操作的顺序，将张量的维度从 (B, C, H, W) 转换为 (B, H, W, C)

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

对比下 python 结果：

可以看到输出数据都相同，这个就是 sample_permute 案例

3.3 sample_reshape

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

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

    auto conv          = network.addConvolutionNd(*data, 3, nvinfer1::DimsHW{3, 3}, mWts["conv.weight"], mWts["conv.bias"]);
    conv->setName("conv1");
    conv->setStride(nvinfer1::DimsHW(1, 1));

    auto reshape       = network.addShuffle(*conv->getOutput(0));
    reshape->setReshapeDimensions(nvinfer1::Dims3{1, 3, -1});
    reshape->setSecondTranspose(nvinfer1::Permutation{0, 2, 1});      
    reshape->setName("reshape + permute1");

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

reshape 操作的实现主要是通过以下几行代码实现的：

auto reshape = network.addShuffle(*conv->getOutput(0));
reshape->setReshapeDimensions(nvinfer1::Dims3{1, 3, -1});
reshape->setSecondTranspose(nvinfer1::Permutation{0, 2, 1});      
reshape->setName("reshape + permute1");

其中：

• network.addShuffle 方法可添加一个 reshape 层，该层也可以执行转置操作
• *conv->getOutput(0) 是卷积层的输出张量，作为 reshape 层的输入
• reshape->setReshapeDimensions(nvinfer1::Dims3{1, 3, -1}) 设置 reshape 的维度为 {1, 3, -1}：
  • 这里的 1 表示批量大小
  • 3 表示输出的通道数
  • -1 表示自动计算该维度的大小，以适应输入和输出的元素总数一致
• reshape->setSecondTranspose(nvinfer1::Permutation{0, 2, 1}) 设置转置操作的顺序，将张量的维度从 (B, C, W) 转换为 (B, W, C)，即将第三维和第二维交换

这个有个点需要大家注意，因为 reshape 和 transpose 都属于 iShuffleLayer 做的事情，所以需要指明是 reshape 在前还是 transpose 在前。另外这里我们可以看到 reshape 和 permute 操作被组合在一个 Shuffle 层中，这种操作可以优化计算效率，是 TensorRT 的一种层融合优化方式

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

对比下 python 结果：

可以看到输出数据都相同，这个就是 sample_reshape 案例

3.4 sample_batchNorm

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

void Model::build_batchNorm(nvinfer1::INetworkDefinition& network, map<string, nvinfer1::Weights> mWts) {
    auto data          = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});
    auto conv          = network.addConvolutionNd(*data, 3, nvinfer1::DimsHW{3, 3}, mWts["conv.weight"], mWts["conv.bias"]);
    conv->setName("conv1");
    conv->setStride(nvinfer1::DimsHW(1, 1));

    float* gamma   = (float*)mWts["norm.weight"].values;
    float* beta    = (float*)mWts["norm.bias"].values;
    float* mean    = (float*)mWts["norm.running_mean"].values;
    float* var     = (float*)mWts["norm.running_var"].values;
    float  eps     = 1e-5;
    
    int    count   = mWts["norm.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 scale = network.addScale(*conv->getOutput(0), nvinfer1::ScaleMode::kCHANNEL, shifts_weights, scales_weights, pows_weights);
    scale->setName("batchNorm1");

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

值得注意的是由于 TensorRT 没有原生的 BatchNorm 层实现，这里用 IScaleLayer 来模拟 BatchNorm 的计算，主要步骤如下：

1. BatchNorm 的参数获取和计算

• 提取 BatchNorm 所需的参数：gamma、beta、mean 和 var，分别对应 BN 的权重、偏置、均值和方差
• eps 是一个小值，防止在计算过程中除零
• 计算 scales、shifts 和 pows，用于在 IScaleLayer 中实现 BatchNorm：
  • scales[i] = gamma[i] / sqrt(var[i] + eps)：计算缩放因子
  • shifts[i] = beta[i] - (mean[i] * gamma[i] / sqrt(var[i] + eps))：计算偏移量
  • pows[i] = 1.0：设置幂次为1，表示不进行额外的幂次操作

float* gamma   = (float*)mWts["norm.weight"].values;
float* beta    = (float*)mWts["norm.bias"].values;
float* mean    = (float*)mWts["norm.running_mean"].values;
float* var     = (float*)mWts["norm.running_var"].values;
float  eps     = 1e-5;

int    count   = mWts["norm.running_var"].count;

float* scales  = (float*)malloc(count * sizeof(float));
float* shifts  = (float*)malloc(count * sizeof(float));
float* pows    = (float*)malloc(count * sizeof(float));

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;
}

2. 创建 Weights 对象

• 将计算得到的 scales、shifts 和 pows 转换为 TensorRT 的 Weights 对象，指定数据类型为浮点型，数量为 count

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};

3. 创建 IScaleLayer 以模拟 BatchNorm

• 使用 network.addScale 方法添加一个 IScaleLayer 层，用于模拟 BatchNorm 的计算
• ScaleMode::kCHANNEL 指定按通道（channel）进行缩放
• 将 shifts_weights、scales_weights 和 pows_weights 作为 IScaleLayer 的参数

auto scale = network.addScale(*conv->getOutput(0), nvinfer1::ScaleMode::kCHANNEL, shifts_weights, scales_weights, pows_weights);
scale->setName("batchNorm1");

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

对比下 python 结果：

可以看到输出数据都相同，这个就是 sample_batchNorm 案例

3.5 sample_cbr

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

void Model::build_cbr(nvinfer1::INetworkDefinition& network, map<string, nvinfer1::Weights> mWts) {
    auto data          = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});
    auto conv          = network.addConvolutionNd(*data, 3, nvinfer1::DimsHW{3, 3}, mWts["conv.weight"], mWts["conv.bias"]);
    conv->setName("conv1");
    conv->setStride(nvinfer1::DimsHW(1, 1));

    float* gamma   = (float*)mWts["norm.weight"].values;
    float* beta    = (float*)mWts["norm.bias"].values;
    float* mean    = (float*)mWts["norm.running_mean"].values;
    float* var     = (float*)mWts["norm.running_var"].values;
    float  eps     = 1e-5;
    
    int    count   = mWts["norm.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(*conv->getOutput(0), nvinfer1::ScaleMode::kCHANNEL, shifts_weights, scales_weights, pows_weights);
    bn->setName("batchNorm1");

    auto leaky = network.addActivation(*bn->getOutput(0), nvinfer1::ActivationType::kLEAKY_RELU);
    leaky->setName("leaky1");

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

和前面的 sample_batchNorm 案例非常像，这边添加了一个激活函数 Leaky_ReLU：

auto leaky = network.addActivation(*bn->getOutput(0), nvinfer1::ActivationType::kLEAKY_RELU);
leaky->setName("leaky1");

其中：

• network.addActivation 方法添加一个激活层
• *bn->getOutput(0) 是 BatchNorm 层的输出张量，作为激活层的输入
• nvinfer1::ActivationType::kLEAKY_RELU 指定激活类型为 Leaky ReLU
• leaky->setName(“leaky1”) 设置激活层的名称为 leaky1

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

对比下 python 结果：

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

Note：最近韩君老师又新增了一些案例，比如 pooling、unsample、deconv 等等，大家感兴趣的可以看看

4. 补充说明

大家如果对 tensorrtx 这个 repo 熟悉的话，会发现 gen_wts.py 将 和这里保存权重的方式一模一样，估计韩君老师也借鉴了这个 repo

gen_wts.py 将 .pt 模型转换为 .wts 模型其实就是这里的把权重按照指定的格式导出，然后在 C++ 上自己去做解析，另外 tensorrtx 这个 repo 也比较有意思，与 tensorRT_Pro 不同的是，它并没有采取 onnxparser 去构建 network，而是像这里讲的一样通过 C++ API 一层层去搭建 network

这样其实需要考验大家对模型的熟练度以及细节的控制，对技能要求高，而且新模型需要自己一个 layer 一个 layer 写 C++ 代码构建，不具有通用性，但是作者也提供了大量场景模型的构建，可以直接使用

总结

本次课程我们主要学习了另外一种构建 network 的方式，与之前利用 onnxparser 解析 onnx 不同，我们这里先将模型的 weights 保存下来然后在 C++ 上去解析构建一个 mapdas，接着通过调用 network 的 addXXX 来构建各个 layer 层，其它的部分和之前的案例没什么区别。值得注意的是这里展示的案例都是非常简单的，真正要构建一个 model 比如 yolo、transformer 等等是比较复杂的，需要将它们封装成一个个模块并做单元测试来验证每个 module 的功能是否正常

OK，以上就是 5.5 小节案例的全部内容了，下节我们来学习 5.6 小节 TensorRT 搭建网络时的模块化思想，敬请期待😄

下载链接

• tensorrt_starter源码
• 5.5-build-model案例文件

参考

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