用C++语言写推理TensorRT代码

下面是一个简单的 C++ TensorRT 推理代码示例：

#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <numeric>
#include <chrono>
#include <cuda_runtime_api.h>
#include <NvInfer.h>
#include <NvInferRuntime.h>

using namespace std;
using namespace nvinfer1;
using namespace nvuffparser;

int main(int argc, char** argv)
{
    // 读取模型文件和权重文件
    const string engine_file_path = "model.engine";
    const string weight_file_path = "model.wts";
    ifstream engine_file(engine_file_path, ios::binary);
    ifstream weight_file(weight_file_path, ios::binary);
    engine_file.seekg(0, engine_file.end);
    const size_t engine_size = engine_file.tellg();
    engine_file.seekg(0, engine_file.beg);
    vector<char> engine_data(engine_size);
    engine_file.read(engine_data.data(), engine_size);
    weight_file.seekg(0, weight_file.end);
    const size_t weight_size = weight_file.tellg();
    weight_file.seekg(0, weight_file.beg);
    vector<char> weight_data(weight_size);
    weight_file.read(weight_data.data(), weight_size);

    // 创建推理引擎
    IRuntime* runtime = createInferRuntime(gLogger);
    ICudaEngine* engine = runtime->deserializeCudaEngine(engine_data.data(), engine_size, nullptr);
    IExecutionContext* context = engine->createExecutionContext();

    // 分配输入和输出内存
    const int input_index = engine->getBindingIndex("input");
    const int output_index = engine->getBindingIndex("output");
    const int batch_size = 1;
    const int input_size = 224 * 224 * 3;
    const int output_size = 1000;
    float* input_data = new float[batch_size * input_size];
    float* output_data = new float[batch_size * output_size];
    cudaMalloc(&input_data, batch_size * input_size * sizeof(float));
    cudaMalloc(&output_data, batch_size * output_size * sizeof(float));

    // 创建 CUDA 流
    cudaStream_t stream;
    cudaStreamCreate(&stream);

    // 创建输入数据
    for (int i = 0; i < batch_size * input_size; ++i) {
        input_data[i] = i % 255 / float(255.0);
    }

    // 执行推理
    context->enqueueV2(&input_data, stream, nullptr);
    cudaStreamSynchronize(stream);

    // 获取输出数据
    cudaMemcpy(output_data, output_data, batch_size * output_size * sizeof(float), cudaMemcpyDeviceToHost);

    // 打印输出结果
    for (int i = 0; i < output_size; ++i) {
        cout << output_data[i] << " ";
    }
    cout << endl;

    // 释放资源
    cudaFree(input_data);
    cudaFree(output_data);
    context->destroy();
    engine->destroy();
    runtime->destroy();
    cudaStreamDestroy(stream);

    return 0;
}

这个示例代码假设已经有了一个 TensorRT 模型引擎文件 `model.engine` 和对应的权重文件 `model.wts`，并且模型的输入名为 `input`，输出名为 `output`。代码中首先读取了模型文件和权重文件，并使用 `createInferRuntime` 函数创建了一个 TensorRT 运行时对象 `runtime`，然后使用 `deserializeCudaEngine` 函数将模型引擎反序列化成 `engine` 对象，使用 `createExecutionContext` 函数创建了一个执行上下文 `context` 对象。

接下来，代码中分配了输入和输出内存，并使用 `getBindingIndex` 函数获取了输入和输出的索引，用 `enqueueV2` 函数执行推理，使用 `cudaMemcpy` 函数将输出数据从 GPU 内存复制到主机内存，并打印了输出结果。

最后，代码中释放了分配的内存，并使用 `destroy` 函数释放了 TensorRT 对象。