tensorRT C++部署实战(yolov5/yolov8/yolov5seg/yolov8seg)

  tensorRT是NVIDIA推出的用于深度学习模型加速的工具库,将Yolov5/Yolov8与TensorRT结合使用，可以在NVIDIA的GPU上进行高效推理。使用C++中使用TensorRT加载和运行Yolov5/Yolov8模型流程主要分为以下三步，

• 导出Yolov5/Yolov8的ONNX格式模型。
• 将ONNX模型序列化，首先实例化Logger，然后创建Builder/Network对象，使用Parser解析ONNX模型构建Network，再设置Config参数优化网络，最终转换为序列化模型，保存为TensorRT引擎。
• 加载引擎并反序列化模型，为输入分配内存，首先拷贝模型输入数据（HostToDevice）执行模型推理，然后拷贝模型输出数据（DeviceToHost）解析结果。

附上代码链接：

完整代码：https://download.csdn.net/download/qq_36801705/89633517
环境配置：https://blog.csdn.net/qq_36801705/article/details/141056249
参考链接：https://mp.weixin.qq.com/s/tWd-o4sRV6EMOAd1b4WFHQ

1.模型导出

python export.py --weights yolov5s.pt --include onnx --imgsz 640 640

2.模型序列化

tensorRT所有的接口都存放在命名空间nvinfer1中，首先要实例化ILogger接口，然后再创建IBuilder对象。

class Logger : public nvinfer1::ILogger {
public:
    explicit Logger(nvinfer1::ILogger::Severity severity =
        nvinfer1::ILogger::Severity::kWARNING)
        : severity_(severity) {}

    void log(nvinfer1::ILogger::Severity severity,
        const char* msg) noexcept override {
        if (severity <= severity_) {
            std::cerr << msg << std::endl;
        }
    }
    nvinfer1::ILogger::Severity severity_;
}gLogger;

创建builder对象后，继续构建模型的网络结构，代码如下：

nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);

  由于ONNX模型是现成的，所以这里采用onnx解析器直接从ONNX模型中解析出模型的网络结构。ONNX解析器接口被封装在头文件NvOnnxParser.h中，命名空间为nvonnxparser。创建ONNX解析器对象并加载模型的代码如下：

nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
// onnx文件解析类
// 将onnx文件解析，并填充rensorRT网络结构
nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);
// 解析onnx文件
parser->parseFromFile(onnx_file_path.c_str(), 2);
for (int i = 0; i < parser->getNbErrors(); ++i) {
    std::cout << "load error: " << parser->getError(i)->desc() << std::endl;
}

  模型解析成功后，需要创建一个IBuilderConfig对象来告诉TensorRT该如何对模型进行优化，这一步可以设置工作空间的最大容量与模型的数据精度，TensorRT默认的数据精度为FP32，还可以根据硬件平台是否支持选择FP16或者INT8：

nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
// 设置最大工作空间大小。
config->setMaxWorkspaceSize(1 << 30);    //2的30次方的大小
// 设置模型输出精度
if (usefp)
    config->setFlag(nvinfer1::BuilderFlag::kFP16);
else
    config->setFlag(nvinfer1::BuilderFlag::kINT8);

另外还可以对动态模型预设定尺寸：

// 动态模型预设尺寸，可根据自己实际情况设置。
if (mInputDims.d[0] < 1)
{
    nvinfer1::Dims minInputSize = Dims4(1, 3, 640, 640);
    nvinfer1::Dims medInputSize = Dims4(1, 3, 640, 640);
    nvinfer1::Dims maxInputSize = Dims4(1, 3, 640, 640);
    
    profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kMIN, minInputSize);
    profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kOPT, medInputSize);
    profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kMAX, maxInputSize);
    
}
else
{
    profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kMIN, mInputDims);
    profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kOPT, mInputDims);
    profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kMAX, mInputDims);
}

  最后则是启动引擎优化模型，优化后的序列化模型被保存到IHostMemory对象中，可以保存到本地磁盘(.engine文件)，方便下次直接加载省去优化的时间。

完整代码：

void onnx_to_engine(std::string onnx_file_path, std::string engine_file_path, bool usefp) {

    // 构建器，获取cuda内核目录以获取最快的实现
    // 用于创建config、network、engine的其他对象的核心类
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger);         

    const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); // 显式批处理
    // 解析onnx网络文件
    // tensorRT模型类
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
    // onnx文件解析类
    // 将onnx文件解析，并填充rensorRT网络结构
    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);
    // 解析onnx文件
    parser->parseFromFile(onnx_file_path.c_str(), 2);
    for (int i = 0; i < parser->getNbErrors(); ++i) {
        std::cout << "load error: " << parser->getError(i)->desc() << std::endl;
    }
    cout << "tensorRT load mask onnx model successfully!!!...\n" << endl;

    // 创建推理引擎
    // 创建生成器配置对象。
    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
    // 设置最大工作空间大小。
    config->setMaxWorkspaceSize(1 << 30);    //2的30次方的大小
    // 设置模型输出精度
    if (usefp)
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
    else
        config->setFlag(nvinfer1::BuilderFlag::kINT8);

    auto profile = builder->createOptimizationProfile();
    assert(network->getNbInputs() == 1);
    nvinfer1::Dims mInputDims = network->getInput(0)->getDimensions();
    auto mInputName = network->getInput(0)->getName();
    assert(mInputDims.nbDims == 4);

    // 动态模型预设尺寸，可根据自己实际情况设置。
    if (mInputDims.d[0] < 1)
    {
        nvinfer1::Dims minInputSize = Dims4(1, 3, 640, 640);
        nvinfer1::Dims medInputSize = Dims4(1, 3, 640, 640);
        nvinfer1::Dims maxInputSize = Dims4(1, 3, 640, 640);

        profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kMIN, minInputSize);
        profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kOPT, medInputSize);
        profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kMAX, maxInputSize);

    }
    else
    {
        profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kMIN, mInputDims);
        profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kOPT, mInputDims);
        profile->setDimensions(mInputName, nvinfer1::OptProfileSelector::kMAX, mInputDims);
    }

    config->addOptimizationProfile(profile);

    // 创建推理引擎
    nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    // 将推理引擎保存到本地
    std::cout << "try to save engine file now....." << std::endl;
    std::ofstream engine_file(engine_file_path, std::ios::binary);
    if (!engine_file) {
        std::cerr << "could not open plan output file" << std::endl;
        return;
    }
    // 将模型转化为文件流数据
    nvinfer1::IHostMemory* engine_stream = engine->serialize();
    // 将文件保存到本地
    engine_file.write(reinterpret_cast<const char*>(engine_stream->data()), engine_stream->size());
    // 销毁创建的对象
    engine_stream->destroy();
    engine->destroy();
    network->destroy();
    parser->destroy();
    std::cout << "convert onnx model to TensorRT engine model successfully!" << std::endl;
}

3.模型推理

模型推理流程如下所示:

首先要从本地加载.engine文件以及反序列化

char* serialized_engine{ nullptr }; //char* serialized_engine==nullptr;  开辟空指针后 要和new配合使用
size_t serialized_size{ 0 };
std::ifstream file(modelPath, std::ios::binary);
if (file.good()) {
    std::cout << "load engine success!" << std::endl;
    file.seekg(0, file.end);
    serialized_size = file.tellg();
    
    file.seekg(0, file.beg);
    serialized_engine = new char[serialized_size];
    file.read(serialized_engine, serialized_size);
    file.close();
}
else {
    std::cout << "load engine failed!" << std::endl;
    std::abort();
}
Logger logger;
// 反序列化引擎
initLibNvInferPlugins(&logger, "");
this->runtime = nvinfer1::createInferRuntime(logger);
// 推理引擎
// 保存模型的模型结构、模型参数以及最优计算kernel配置；
// 不能跨平台和跨TensorRT版本移植
this->engine = runtime->deserializeCudaEngine(serialized_engine, serialized_size);

  然后通过createExecutionContext()函数创建一个IExecutionContext对象来管理推理的过程。由于模型的推理是在GPU上进行的，所以会存在搬运输入、输出数据的操作，因此有必要在GPU上创建内存区域用于存放输入、输出数据。模型输入、输出的尺寸可以通过ICudaEngine对象的接口获取，根据这些信息我们可以先为模型分配输入、输出缓存区：
// 由engine创建，可创建多个对象，进行多推理任务

this->context = engine->createExecutionContext();

int kInputIndex = engine->getBindingIndex(kInputNodeName);
auto kIutputDims = engine->getBindingDimensions(kInputIndex);
this->kInputSize = 1;

this->kInputShape = {std::max((int)kIutputDims.d[0], 1), (int)kIutputDims.d[1], (int)kIutputDims.d[2], (int)kIutputDims.d[3] };


for(int i=0; i< kInputShape.size();i++)
{    
    kInputSize *= kInputShape[i];
}

kOutputIndexDet = engine->getBindingIndex(this->kOutputNodeDet);
auto kOutputDimsDet = engine->getBindingDimensions(kOutputIndexDet);
kOutputShapeDet = { std::max((int)kOutputDimsDet.d[0], 1), (int)kOutputDimsDet.d[1], (int)kOutputDimsDet.d[2]};
this-> kOutputSizeDet = 1;
for (int i = 0; i < kOutputShapeDet.size(); i++)
{
    kOutputSizeDet *= kOutputShapeDet[i];
}
由于seg和det任务的输出维度是不一样的，同一种任务yoloV8和yoloV5的输出的维度也存在区别，以检测任务为例，V5的输出维度[1,25200,85]，V8维度为[1,84,8400]，因此可以根据输出维度判断模型和任务的类别：
this->yolov8 = kOutputShapeDet[1] < kOutputShapeDet[2] ? true : false;
this->netWidth = kOutputShapeDet[1] < kOutputShapeDet[2] ? kOutputShapeDet[1] : kOutputShapeDet[2];
int numOutputs = engine->getNbBindings();
if (numOutputs > 2)
{
    this->instSeg = true;
    std::cout << "Instance Segmentation" << std::endl;
    kOutputIndexSeg = engine->getBindingIndex(this->kOutputNodeSeg);
    auto kOutputIDimsSeg = engine->getBindingDimensions(kOutputIndexSeg);
    vector<int>kOutputIShapeSeg = { std::max((int)kOutputIDimsSeg.d[0], 1), (int)kOutputIDimsSeg.d[1], (int)kOutputIDimsSeg.d[2], (int)kOutputIDimsSeg.d[3] };
    
    /*kOutputIShapeSeg[2] = kOutputIShapeSeg[2] <= 1 ? this->widthSeg : kOutputIShapeSeg[2];
    kOutputIShapeSeg[3] = kOutputIShapeSeg[3] <= 1 ? this->heightSeg : kOutputIShapeSeg[3];*/
    
    this->kOutputSizeSeg = 1;
    for (int i = 0; i < kOutputIShapeSeg.size(); i++)
    {
        kOutputSizeSeg *= kOutputIShapeSeg[i];
    }
}
else
    std::cout << "Object Detection" << std::endl;

  根据预设的大小申请输入输出内存，调用IExecutionContext对象的enqueueV2()函数进行异步地推理操作。模型推理成功后，其输出数据被拷贝到buffer中然后按照对应的输出数据排布规则解析即可。

void Yolo::ainfer(std::vector<float>& input, float* outputs, float* outputsSeg)
{

    void* buffers[3];

    // 创建GPU显存输出缓冲区
    CheckCuda(cudaMalloc(&buffers[kInputIndex], batchSize * kInputSize * sizeof(float)));
    CheckCuda(cudaMalloc(&buffers[kOutputIndexDet], batchSize * kOutputSizeDet * sizeof(float)));
    if (this->instSeg)
        CheckCuda(cudaMalloc(&buffers[kOutputIndexSeg], batchSize * kOutputSizeSeg * sizeof(float)));

    // 创建输入cuda流
    cudaStream_t stream;
    CheckCuda(cudaStreamCreate(&stream));

    // 输入数据由内存到GPU显存
    CheckCuda(cudaMemcpyAsync(buffers[kInputIndex], input.data(), batchSize * kInputSize * sizeof(float), cudaMemcpyHostToDevice, stream));

    // 模型推理
    this->context->enqueueV2(buffers, stream, nullptr);
    // 将GPU数据同步到CPU中
    CheckCuda(cudaMemcpyAsync(outputs, buffers[kOutputIndexDet], batchSize * kOutputSizeDet * sizeof(float), cudaMemcpyDeviceToHost, stream));
    if (this->instSeg)
        CheckCuda(cudaMemcpyAsync(outputsSeg, buffers[kOutputIndexSeg], batchSize * kOutputSizeSeg * sizeof(float), cudaMemcpyDeviceToHost, stream));
    cudaStreamSynchronize(stream);

    // Release stream and buffers
    cudaStreamDestroy(stream);
    CheckCuda(cudaFree(buffers[kInputIndex]));
    CheckCuda(cudaFree(buffers[kOutputIndexDet]));
    if (this->instSeg)
        CheckCuda(cudaFree(buffers[kOutputIndexSeg]));

}