快捷導航

使用C++將yolov8 onnx格式轉(zhuǎn)化為tensorrt格式

更新時間：2024年11月15日 09:31:22 作者：愛喝生椰的程序員

這篇文章主要為大家詳細介紹了如何使用C++將yolov8 onnx格式轉(zhuǎn)化為tensorrt格式,文中的示例代碼講解詳細,感興趣的小伙伴可以了解一下

我們使用TensorRT進行加速推理時，需要先將onnx格式轉(zhuǎn)化為tensorrt格式，以下是使用C++來進行轉(zhuǎn)化代碼以及對應(yīng)的CMakeLists.txt文件

操作系統(tǒng)：ubuntu20.04

C++代碼：

// main.cpp
#include <iostream>
#include <memory>
#include <fstream>
#include <assert.h>
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "common.h"

class Logger : public nvinfer1::ILogger{
    void log(Severity severity, const char* msg) noexcept override {
        // suppress info-level messages
        if (severity <= Severity::kWARNING)
            std::cout << msg << std::endl;
    }
} logger;


int main(int argc, char** argv){
    if(argc !=2){
        std::cerr << "usage: ./build [onnx_file_path]" <<std::endl;
        return -1;
    }

    // 獲取onnx文件路徑
    char* onnx_file_path = argv[1];

    //================1.創(chuàng)建builder====================
    auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
    if(!builder){
        std::cerr << "Failed to creater builder" <<std::endl;
        return -1;
    }

    //===============2.創(chuàng)建network=====================
    const auto explicitBatch= 1U <<static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));
    if (!network){
        std::cout << "Failed to create network" << std::endl;
        return -1;
    }

    // =============創(chuàng)建onnxparser用于解析onnx文件===========
    auto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, logger));
    // 調(diào)用onnxparser的parseFromFile方法解析onnx文件
    auto parsed = parser->parseFromFile(onnx_file_path, static_cast<int>(nvinfer1::ILogger::Severity::kWARNING));
    if (!parsed){
        std::cout << "Failed to parse onnx file" << std::endl;
        return -1;
    }
    // 配置網(wǎng)絡(luò)參數(shù)
    auto input = network->getInput(0);
    auto profile = builder->createOptimizationProfile();
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4{1, 3, 960, 960}); // 設(shè)置最小尺寸
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4{1, 3, 960, 960}); // 設(shè)置最優(yōu)尺寸
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4{1, 3, 960, 960}); // 設(shè)置最大尺寸

    //==============創(chuàng)建config配置=====================
    auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    if (!config){
        std::cout << "Failed to create config" << std::endl;
        return -1;
    }

    config->addOptimizationProfile(profile);
    // 設(shè)置精度，設(shè)置為FP16，設(shè)置為INT8需要額外calibrator 
    config->setFlag(nvinfer1::BuilderFlag::kFP16);
    // 設(shè)置最大batchsize
    builder->setMaxBatchSize(1);
    // 設(shè)置最大工作空間
    config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 30);

    //創(chuàng)建流
    auto profileStream = samplesCommon::makeCudaStream();
    if(!profileStream){
        return -1;
    }
    config->setProfileStream(*profileStream);

    // ==========創(chuàng)建engine ==========
    auto plan = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));
    if (!plan){
        std::cout << "Failed to create engine" << std::endl;
        return -1;
    }

    // ========== 5. 序列化保存engine ==========
    std::ofstream engine_file("./yolov8x.engine", std::ios::binary);
    assert(engine_file.is_open() && "Failed to open engine file");
    engine_file.write((char *)plan->data(), plan->size());
    engine_file.close();

    std::cout << "Engine build success!" << std::endl;
    return 0;
}

CMakeLists.txt文件：

cmake_minimum_required(VERSION 3.10)
project(TensorRT_Test LANGUAGES CXX CUDA)

set(CMAKE_CUDA_STANDARD 14)
set(CMAKE_CXX_STANDARD 14)

# 添加頭文件路徑 cuda tensorRT
include_directories(/usr/local/cuda-11.8/include)
include_directories(/xxx/tensorRT/TensorRT-8.6.1.6/include)
include_directories(/xxx/tensorRT/TensorRT-8.6.1.6/samples/common/)

# 添加庫文件
link_directories(/usr/local/cuda-11.8/lib64)
link_directories(/xxx/tensorRT/TensorRT-8.6.1.6/lib)

add_executable(build main.cpp)
target_link_libraries(build nvinfer nvonnxparser cudart)

注意：

在CMakeLists.txt中cuda及TensorRT的頭文件和庫文件路徑需要改成自己的

使用方法：

以yolov8為例，使用cmake編譯好后會生成build可執(zhí)行文件，執(zhí)行以下命令即可等待生成yolov8x.engine文件

./build <onnx_path>

到此這篇關(guān)于使用C++將yolov8 onnx格式轉(zhuǎn)化為tensorrt格式的文章就介紹到這了,更多相關(guān)C++ yolov8 onnx轉(zhuǎn)tensorrt內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: