TensorRT学习系列一：使用自定义网络构建和运行TensorRT模型

---

一、前言

随着深度学习技术的发展，许多实际应用都采用了神经网络模型来处理复杂任务，如图像识别、语音识别、自然语言处理等。然而，在部署这些模型时，如何实现在生产环境中高效且准确地执行推理是一个重要问题。这就需要利用专门针对推理优化的工具和技术，其中NVIDIA的TensorRT就是这样一款高性能的深度学习推理优化器。

TensorRT是一个由NVIDIA开发的高性能推理平台，旨在加速深度学习模型的推断阶段，并最大限度地提高GPU的利用率。它通过一系列优化技术，包括模型剪枝、量化以及算法选择等，能够在保持模型精度的同时，显著提升模型在生产环境中的运行效率。
TensorRT广泛应用于各种场景，例如自动驾驶、医疗影像分析、云计算服务、实时推荐系统等，特别是在需要实时响应或处理大量数据流的情况下，其优势尤为明显。

本文是“TensorRT学习系列”之一，我们将专注于如何使用自定义网络构建和运行TensorRT模型。我们将详细介绍如何将训练好的模型导入TensorRT，并对其进行优化以实现高效的推理。此外，本文还将特别关注TensorRT构建和运行时的新旧API对比说明，帮助读者更好地理解TensorRT的工作原理，以便在实际项目中做出最佳实践的选择。

在这个过程中，我们不仅会探讨TensorRT的基本使用方法，还会涉及到性能优化策略和常见问题的解决方案。希望通过本文的学习，你能掌握如何有效地利用TensorRT提升深度学习模型的推理性能。

关于运行环境的搭建和TensorRT的安装，请参考我另一篇博客：Ubuntu下安装ONNX、ONNX-TensorRT、Protobuf和TensorRT，本文不再详述。

---

二、使用TensorRT构建(build)引擎

我们使用的TensorRT 8.6-GA版，比这之前的统称旧版，部分代码可能在这之前的某个版本声明将要废弃(deprecated)，并不一定是这一版开始声明将要废弃的。这里不单独声明，具体可以查询TensorRT的官方api文档。另外，这些声明将要废弃的代码在这一版还是可以运行的。
下面是TensorRT build engine的主要过程

1. 创建logger
2. 创建builder
3. 创建network (builder->network)
4. 配置参数 (builder->config)
5. 生成engine和序列化保存 (builder->serialized_engine(network, config))
6. 释放资源 (delete)

我们分别进行代码说明：

说明：我们在代码中会有#define oldversion 0，这个oldversion的宏就是指代之前的版本写法。
另外，下面很多代码里有缩进，那是因为本身是在main函数里。

1. 创建logger

这个写法相对固定，主要是创建builder的时候需要，所以写在前面

class Logger : public nvinfer1::ILogger {
    void log(Severity severity, const char *msg) noexcept override {
        if (severity != Severity::kINFO) {
            std::cout << msg << std::endl;
        }
    }
} gLogger;

2. 创建builder

这个写法相对固定，新旧版本之间没有太大差异。

	nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(gLogger);

3. 创建network

这是一种指定显性batch的写法，也可以直接指定参数为1。注意使用的是V2。

	auto explictBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);    
	nvinfer1::INetworkDefinition *network = builder->createNetworkV2(explictBatch);

我们讲自定义网络，手动添加输入层，全连接层和激活层作为案例说明。

    const int input_size = 3;
    const int output_size = 2;
    // 标记输入层和创建输入数据
    nvinfer1::ITensor *input = network->addInput("data", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, input_size, 1, 1});
    const float *fc1_weight_data = new float[input_size * output_size]{0.1,0.2,0.3,0.4,0.5,0.6};
    const float *fc1_bias_data = new float[2]{0.1,0.5};
    nvinfer1::Weights fc1_weight{nvinfer1::DataType::kFLOAT, fc1_weight_data, input_size * output_size}; 
    nvinfer1::Weights fc1_bias{nvinfer1::DataType::kFLOAT, fc1_bias_data, output_size};
// 下面是比较核心的部分，是新旧版本目前看来差异是大的部分
#if oldversion
    // 设定最大batch size
    builder->setMaxBatchSize(1);
    nvinfer1::IFullyConnectedLayer *fc1 = network->addFullyConnected(*input, output_size, fc1_weight, fc1_bias);
    nvinfer1::IActivationLayer *sigmoid = network->addActivation(*fc1->getOutput(0), nvinfer1::ActivationType::kSIGMOID); 
#else
    // 将输入张量转换为2D，以进行矩阵乘法
    nvinfer1::IShuffleLayer* shuffleLayer = network->addShuffle(*input);
    shuffleLayer->setReshapeDimensions(nvinfer1::Dims2{input_size, 1});
    // 创建权重矩阵的常量层
    nvinfer1::IConstantLayer* weightLayer = network->addConstant(nvinfer1::Dims2{output_size, input_size}, fc1_weight);
    // 添加矩阵乘法层
    nvinfer1::IMatrixMultiplyLayer* matMulLayer = network->addMatrixMultiply(
        *weightLayer->getOutput(0), nvinfer1::MatrixOperation::kNONE,
        *shuffleLayer->getOutput(0), nvinfer1::MatrixOperation::kNONE
    );
    // 添加偏置
    nvinfer1::IConstantLayer* biasLayer = network->addConstant(nvinfer1::Dims2{output_size, 1}, fc1_bias);
    nvinfer1::IElementWiseLayer* addBiasLayer = network->addElementWise(
        *matMulLayer->getOutput(0), *biasLayer->getOutput(0), nvinfer1::ElementWiseOperation::kSUM);
    // 添加激活层
    nvinfer1::IActivationLayer* sigmoid = network->addActivation(*addBiasLayer->getOutput(0), nvinfer1::ActivationType::kSIGMOID);
#endif
		// 标记输出层
    sigmoid->getOutput(0)->setName("output");
    network->markOutput(*sigmoid->getOutput(0));

创建全连接层和激活层的方法在新旧版本中差异较大，需要多关注。

4. 配置参数

createBuilderConfig接口被用来指定TensorRT应该如何优化模型。这里应该注意，在旧版builder->setMaxBatchSize(1);这个代码一定要有，否则运算结果不对。设置最大工作区大小，相当于设置内存池，新旧版本api不同。

#if oldversion
    builder->setMaxBatchSize(1);
#endif
    nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();
// 设置最大工作区大小，这个就是api的不同
#if oldversion
    config->setMaxWorkspaceSize(1 << 28);
#else
    config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1<<28);
#endif

5. 构建并序列化引擎并保存

这个在旧版本中是分两步，构建和序列化是分开的，在新版中合为一个api。这所以要序列化，主要是构建过程比较长，一般序列化后要保存成文件，这里不作过多介绍。

// 构建并序列化引擎
#if oldversion
    nvinfer1::ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
    if (!engine) {
        std::cerr << "build engine failed" << std::endl;
		// 将之前堆内存分配的变量释放
        return -1;
    }
    nvinfer1::IHostMemory *serialized_engine = engine->serialize();
#else
// 使用 buildSerializedNetwork 构建并序列化网络
    nvinfer1::IHostMemory* serialized_engine = builder->buildSerializedNetwork(*network, *config);
    if (!serialized_engine) {
        std::cerr << "build engine failed" << std::endl;
		// 将之前堆内存分配的变量释放
        return -1;
    }
#endif

6. 释放资源

这块比较简单，不详述。就是注意一下，在旧版本中使用对象的destory方法，新版本直接delete。再一个就是注意按照构造的顺序反过来进行释放。

#if oldversion
	...
	builder.destroy();
#else
	...
	delete builder;
#endif

---

三、使用TensorRT运行(runtime)模型

上面构建好了引擎，我们需要运行模型进行推理，下面是runtime的主要过程

1. 创建一个runtime对象
2. 反序列化生成engine
3. 创建一个执行上下文
4. 填充数据
5. 执行推理
6. 释放资源
   我们分别进行代码说明：

logger和上部分一样，不单独列出来了
下面我部分采取了智能指针的写法，自动释放资源。

1. 创建一个runtime对象

标准写法。

	// 创建一个runtime对象
	auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(gLogger));

2. 反序列化生成engine

loadData是自己实现从文件读取序列化数据，不再详述。

    // 反序列化生成engine
    std::string serialized_engine_file = argv[1];
    std::vector<char> engine_data = loadData(serialized_engine_file);
    auto engine = std::unique_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(engine_data.data(), engine_data.size()));
    if (engine == nullptr) {
        std::cout << "deserializeCudaEngine failed" << std::endl;
        // 因为采用的是智能指针，不必再手动释放资源。
        return -1;
    }

3. 创建一个执行上下文

    auto context = std::unique_ptr<nvinfer1::IExecutionContext>(engine->createExecutionContext());
    if (context == nullptr) {
        std::cout << "createExecutionContext failed" << std::endl;
        return -1;
    }
#endif

4. 填充数据

这块是cuda编程非常常规的写法，这里不再详细说明。

    // 输入数据
    // float host_input_data = new float[input_size]{1,2,3};
    std::unique_ptr<float[]> host_input_data(new float[input_size]{2,4,8});
    int host_intput_size = input_size * sizeof(float);
    float *device_input_data = nullptr;
    // 输出数据
    // float host_output_data = new float[output_size]{0};
    std::unique_ptr<float[]> host_output_data(new float[output_size]{0});
    int host_output_size = output_size * sizeof(float);
    float *device_output_data = nullptr;
    // 创建 CUDA 流
    cudaStream_t stream;
    cudaStreamCreate(&stream);
    // 申请device内存
    cudaMalloc((void **)&device_input_data, host_intput_size);
    cudaMalloc((void**)&device_output_data, host_output_size);
    cudaMemcpyAsync(device_input_data, host_input_data.get(), host_intput_size, cudaMemcpyHostToDevice, stream);

5. 执行推理

我们注意到在旧版本中使用的是enqueueV2，但在新版本中用的是executeV2，而不是enqueueV3。根据网上资料enqueueV3和executeV2作用是相同的，但是只接受一个stream作为参数，但没有实践成功。我一直想研究一个enqueueV3，因为我开始感觉这是最正规的做法，后来看了TensorRT官方sample里面的案例，用的就是executeV2。所以executeV2才是目前最正规的做法。

    // 准备绑定缓冲区
    void* bindings[] = {device_input_data, device_output_data};
#if oldversion
    bool status = context->enqueueV2((void**)bindings, stream, nullptr);
#else
    bool status = context->executeV2(bindings);
#endif

6. 释放资源

略

---

四、总结

本文深入探讨了如何使用自定义网络构建和运行TensorRT模型。首先，介绍了TensorRT的基础知识，包括其基本概念和应用场景。接着，详细讲解了如何将自定义网络转换为TensorRT模型，并在此过程中进行了新旧API的对比说明。随后，展示了如何使用TensorRT运行模型并获取结果。在这些过程中， 列举了一些常见的问题及其解决方案，希望能帮助读者在实际操作中避免这些问题。

五、参考

我在github上有具体的实现，希望给读者以参考。后续的关于TensorRT的分享也会有部分在这个项目中实现。
项目github