随着深度学习技术的发展，许多实际应用都采用了神经网络模型来处理复杂任务，如图像识别、语音识别、自然语言处理等。然而，在部署这些模型时，如何实现在生产环境中高效且准确地执行推理是一个重要问题。这就需要利用专门针对推理优化的工具和技术，其中NVIDIA的TensorRT就是这样一款高性能的深度学习推理优化器。
TensorRT是一个由NVIDIA开发的高性能推理平台，旨在加速深度学习模型的推断阶段，并最大限度地提高GPU的利用率。它通过一系列优化技术，包括模型剪枝、量化以及算法选择等，能够在保持模型精度的同时，显著提升模型在生产环境中的运行效率。
TensorRT广泛应用于各种场景，例如自动驾驶、医疗影像分析、云计算服务、实时推荐系统等，特别是在需要实时响应或处理大量数据流的情况下，其优势尤为明显。
本文是“TensorRT学习系列”之一，我们将专注于如何使用自定义网络构建和运行TensorRT模型。我们将详细介绍如何将训练好的模型导入TensorRT，并对其进行优化以实现高效的推理。此外，本文还将特别关注TensorRT构建和运行时的新旧API对比说明，帮助读者更好地理解TensorRT的工作原理，以便在实际项目中做出最佳实践的选择。
在这个过程中，我们不仅会探讨TensorRT的基本使用方法，还会涉及到性能优化策略和常见问题的解决方案。希望通过本文的学习，你能掌握如何有效地利用TensorRT提升深度学习模型的推理性能。
关于运行环境的搭建和TensorRT的安装，请参考我另一篇博客：Ubuntu下安装ONNX、ONNX-TensorRT、Protobuf和TensorRT，本文不再详述。
我们使用的TensorRT 8.6-GA版，比这之前的统称旧版，部分代码可能在这之前的某个版本声明将要废弃(deprecated)，并不一定是这一版开始声明将要废弃的。这里不单独声明，具体可以查询TensorRT的官方api文档。另外，这些声明将要废弃的代码在这一版还是可以运行的。
下面是TensorRT build engine的主要过程
我们分别进行代码说明：
说明：我们在代码中会有
#define oldversion 0，这个oldversion的宏就是指代之前的版本写法。
另外，下面很多代码里有缩进，那是因为本身是在main函数里。
这个写法相对固定，主要是创建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;
这个写法相对固定，新旧版本之间没有太大差异。
nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(gLogger);
这是一种指定显性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));
创建全连接层和激活层的方法在新旧版本中差异较大，需要多关注。
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
这个在旧版本中是分两步，构建和序列化是分开的，在新版中合为一个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
这块比较简单，不详述。就是注意一下，在旧版本中使用对象的destory方法，新版本直接delete。再一个就是注意按照构造的顺序反过来进行释放。
#if oldversion
...
builder.destroy();
#else
...
delete builder;
#endif
上面构建好了引擎，我们需要运行模型进行推理，下面是runtime的主要过程
logger和上部分一样，不单独列出来了
下面我部分采取了智能指针的写法，自动释放资源。
标准写法。
// 创建一个runtime对象
auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(gLogger));
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;
}
auto context = std::unique_ptr<nvinfer1::IExecutionContext>(engine->createExecutionContext());
if (context == nullptr) {
std::cout << "createExecutionContext failed" << std::endl;
return -1;
}
#endif
这块是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);
我们注意到在旧版本中使用的是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
略
本文深入探讨了如何使用自定义网络构建和运行TensorRT模型。首先，介绍了TensorRT的基础知识，包括其基本概念和应用场景。接着，详细讲解了如何将自定义网络转换为TensorRT模型，并在此过程中进行了新旧API的对比说明。随后，展示了如何使用TensorRT运行模型并获取结果。在这些过程中， 列举了一些常见的问题及其解决方案，希望能帮助读者在实际操作中避免这些问题。
我在github上有具体的实现，希望给读者以参考。后续的关于TensorRT的分享也会有部分在这个项目中实现。
项目github