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

news2025/1/21 17:59:52

目录

    • 前言
    • 0. 简述
    • 1. 案例运行
    • 2. 代码分析
      • 2.1 main.cpp
      • 2.2 model.cpp
    • 3. 案例
      • 3.1 sample_conv
      • 3.2 sample_permute
      • 3.3 sample_reshape
      • 3.4 sample_batchNorm
      • 3.5 sample_cbr
    • 4. 补充说明
    • 总结
    • 下载链接
    • 参考

前言

自动驾驶之心推出的 《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

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：http://www.coloradmin.cn/o/1976178.html

如若内容造成侵权/违法违规/事实不符，请联系多彩编程网进行投诉反馈，一经查实，立即删除！

相关文章

《学会 SpringMVC 系列 · 写入拦截器 ResponseBodyAdvice》

《学会 SpringMVC 系列 · 写入拦截器 ResponseBodyAdvice》

&#x1f4e2; 大家好&#xff0c;我是 【战神刘玉栋】&#xff0c;有10多年的研发经验&#xff0c;致力于前后端技术栈的知识沉淀和传播。 &#x1f497; &#x1f33b; CSDN入驻不久&#xff0c;希望大家多多支持&#xff0c;后续会继续提升文章质量&#xff0c;绝不滥竽充数…
阅读更多...
3.4数组和特殊矩阵

3.4数组和特殊矩阵

3.4.1数组的定义 数组是由n个相同类型的数据元素构成的有序序列 数组是线性表的推广,一个数组可以视为一个线性表 数组一旦被定义,其长度不会再改变,所以数组只会有存取元素和修改元素的操作 3.4.2数组的存储结构 多维数组 有两种映射方法:按行优先和按列优先 按行优先 …
阅读更多...
2024 年最值得阅读的 10 个外国技术网站

2024 年最值得阅读的 10 个外国技术网站

从网络上数以千计的博客中挑选出最好的技术网站&#xff0c;并根据相关性、权威性、社交媒体关注者和新鲜度进行排名。 1. TechCrunch TechCrunch 是一家领先的科技媒体&#xff0c;致力于深入分析初创公司、评论新的互联网产品和发布科技新闻。该网站是科技专业人士和爱好者…
阅读更多...
【传知代码】实体关系抽取（论文复现）

【传知代码】实体关系抽取（论文复现）

当谈论信息提取领域的最前沿时&#xff0c;实体关系抽取无疑是其中一颗耀眼的明星。从大数据时代的信息海洋中提炼出有意义的关系&#xff0c;不仅是科技进步的体现&#xff0c;更是人类对知识管理和智能决策迫切需求的响应。本文将探索实体关系抽取的核心技术、应用场景及其在…
阅读更多...
域控搭建（windows 2012 R2和win10）

域控搭建（windows 2012 R2和win10）

域控搭建 环境准备 两台windows虚拟机 主域控为&#xff1a;windows server2012 子域为&#xff1a;win10 虚拟机设置网段 Win10网络设置 Windows server2012网络设置 Windows server2012网络适配器 设置 识别成功 更改计算机名字 等待重启 Win10网络适配器 设置 识别成功 …
阅读更多...
opencv-图像透视变换

opencv-图像透视变换

透射变换是视角变化的结果&#xff0c;是指利用透视中心&#xff0c;像点&#xff0c;目标点共线的条件&#xff0c;按透视旋转定律使承影面(透视面)绕迹线(透视轴旋转某一角度&#xff0c;破坏原有的投影光束&#xff0c;仍能保持承影面上投影几何图形不变的变化) 它的本质将图…
阅读更多...
QT实现步进电机控制和IMU数据读取显示

QT实现步进电机控制和IMU数据读取显示

实现功能&#xff1a; 1.两步进电机分别使能和循环运动&#xff0c;可以设置循环次数、循环里分别运行的角度、旋转的速度和加减速度等等&#xff0c;在最下方的表格里显示发送和接收的CAN报文 2.读取水平电机当前位置和速度并画图显示&#xff0c;示波器暂停、缩放、滑动等功…
阅读更多...
CVPR24《Neural Markov Random Field for Stereo Matching》

CVPR24《Neural Markov Random Field for Stereo Matching》

论文地址&#xff1a; https://arxiv.org/abs/2403.11193 源码地址&#xff1a; https://github.com/aeolusguan/NMRF 概述 手工设计的MRF模型在传统的立体匹配中占据主导地位&#xff0c;但与端到端的深度学习模型相比&#xff0c;其建模准确性不足。尽管深度学习大大改进了MR…
阅读更多...
力扣SQL50 修复表中的名字 字符串函数

力扣SQL50 修复表中的名字 字符串函数

Problem: 1667. 修复表中的名字 &#x1f468;‍&#x1f3eb; 参考题解 select user_id, CONCAT(UPPER(left(name, 1)), LOWER(RIGHT(name, length(name) - 1))) as name from Users order by user_id
阅读更多...
SQL注入实例（sqli-labs/less-2）

SQL注入实例（sqli-labs/less-2）

0、初始网页 1、闭合方式判断 当没有闭合符号进行注释时&#xff0c;网页并没有报错&#xff0c;所以可以确定无闭合符号&#xff0c;为数值型注入 2、确定查询表的列数 可以确定列数小于4 ?id1 order by 4 -- 确定查询表的列数为3列 ?id1 order by 3 -- 3、确定回显位置…
阅读更多...
MySQL系列之--关系型数据库以及SQL语句分类之DDL数据库和表的操作

MySQL系列之--关系型数据库以及SQL语句分类之DDL数据库和表的操作

文章目录 前言关系型数据库&#xff08;RDBMS&#xff09;关系型数据库的特点 MySQL数据模型SQL介绍基本语法规则SQL语句的分类DDL的介绍DDL的数据库操作DDL的表操作 前言 上一节MySQL系列之–详细安装教程和启动方法中介绍了MySQL如何安装&#xff0c;以及如何启动和客户端连接…
阅读更多...
c++| c++11左右值引用，完美转发，可变参数模板，functional包装器，bind函数

c++| c++11左右值引用，完美转发，可变参数模板，functional包装器，bind函数

c| c11的新特性 左&#xff0c;右值引用什么是左值&#xff0c;右值左值引用和右值引用右值引用解决什么问题呢&#xff1f;移动构造万能引用形式 完美转发格式 lambada表达式格式 可变参数模板可变参数模板实现打印不同类型emplace_push以list的emplace_back的实现举例包装器b…
阅读更多...
新160个crackme - 020-cosh.3

新160个crackme - 020-cosh.3

运行分析 老规矩&#xff0c;需要破解Name和Serial PE分析 c程序&#xff0c;32位&#xff0c;无壳 静态分析&动态调试 ida查找关键字符串 分析关键函数&#xff0c;得出以下结论&#xff1a;Name、Serial每一位进行亦或计算&#xff0c;若计算结果相等则弹窗成功 算法分析…
阅读更多...
吴恩达机器学习作业-ex7（主成分分析）

吴恩达机器学习作业-ex7（主成分分析）

data1 导入库&#xff0c;读取数据&#xff0c;并进行可视化数据 import numpy as np import scipy.io as sio import matplotlib.pyplot as plt#读取数据 path "./ex7data1.mat" data sio.loadmat(path) # print(data.keys()) X data.get("X") # pri…
阅读更多...
Python数据分析案例58——热门游戏数据分析及其可视化

Python数据分析案例58——热门游戏数据分析及其可视化

案例背景 有哪个男生不喜欢玩游戏呢&#xff1f;就算上了班儿也要研究一下游戏以及热门的游戏。正好这里有个热门的游戏数据集&#xff0c;全球热门游戏数据集来做一下一些可视化的分析。 数据介绍 该文件包含一个数据集&#xff0c;详细说明了多个平台上的各种流行游戏。每个…
阅读更多...
【Golang 面试 - 进阶题】每日 3 题（十五）

【Golang 面试 - 进阶题】每日 3 题（十五）

✍个人博客&#xff1a;Pandaconda-CSDN博客 &#x1f4e3;专栏地址&#xff1a;http://t.csdnimg.cn/UWz06 &#x1f4da;专栏简介&#xff1a;在这个专栏中&#xff0c;我将会分享 Golang 面试中常见的面试题给大家~ ❤️如果有收获的话&#xff0c;欢迎点赞&#x1f44d;收藏…
阅读更多...
Cesium 高德地图暗黑化

Cesium 高德地图暗黑化

Cesium 高德地图暗黑化 高德电子地图实现暗黑效果 // 设置图层滤镜new ImageryLayerFilter({viewer: viewer,imageryLayer: imageryLayer,});
阅读更多...
YOLOv5轻量化改进 | backbone | 结合MobileNetV4（包含多个结构和使用方式）

YOLOv5轻量化改进 | backbone | 结合MobileNetV4（包含多个结构和使用方式）

YOLOv5轻量化改进 | backbone | 结合MobileNetV4（包含多个结构） 本文介绍论文原理介绍网络代码多种yaml设置网络测试及实验结果<!-- 这里放入论文图片 --> &emsp;;本文介绍 本文给大家带来的改进机制是结合MobileNetV4骨干网络，其中来自2024.5月发布的MobileNetV4…
阅读更多...
Pageadmin 漏洞教程

Pageadmin 漏洞教程

上传文件解压拿webshell 用哥斯拉&#xff0c;生成一个asp木马 这个就是我们生成的asp木马 不能直接上传 压缩为zip文件 上传 上传以后我们点击解压 就成功了 然后我们去访问 然后用哥斯拉访问 然后再点击添加 这样就成功了 成功进入
阅读更多...
VSCode在windows系统下的配置简单版

VSCode在windows系统下的配置简单版

参考链接 从零开始的vscode安装及环境配置教程(C/C)(Windows系统)_vscode搭建编译器环境-CSDN博客 vscode生成tasks.json、launch.json、c_cpp_properties.json文件_vscode生成launch.json-CSDN博客 自动生成配置文件简单方便&#xff01;&#xff01;&#xff01; 运行c代…
阅读更多...
推荐文章
最新文章

Copyright @ 2022~2023