linux下 yolov8 tensorrt模型部署

warning: 这篇文章距离上次修改已过205天,其中的内容可能已经有所变动。

在Linux下将YOLOv8模型部署到TensorRT需要以下步骤:

  1. 确保你已经安装了NVIDIA的TensorRT库。
  2. 使用ONNX将YOLOv8模型从PyTorch转换为ONNX格式。
  3. 从ONNX将模型转换为TensorRT的engine文件。
  4. 编写C++代码以使用TensorRT运行模型。

以下是一个简化的示例流程:




# 安装TensorRT
sudo apt-get install tensorrt
 
# 安装ONNX
pip install onnx



# 导入必要的库
import torch
import onnx
 
# 加载YOLOv8模型
model = torch.hub.load('ultralytics/yolov8', 'custom', path='yolov8.pt', source='local')
 
# 将PyTorch模型转换为ONNX格式
torch.onnx.export(model, dummy_input, 'model.onnx', verbose=True, input_names=['input'], output_names=['output'])
 
# 这里的dummy_input是模型需要的任意输入数据



# 使用TensorRT的python API将ONNX模型转换为TensorRT engine
# 需要TensorRT的python API支持
 
from tensorrt.python import trt
 
def convert_onnx_to_trt_model(onnx_model_path, engine_path):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
        builder.max_workspace_size = 1 << 30
        with open(onnx_model_path, 'rb') as model:
            if not parser.parse(model.read()):
                print('ERROR: Failed to parse the ONNX file.')
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                return None
 
        engine = builder.build_cuda_engine(network)
        with open(engine_path, 'wb') as f:
            f.write(engine.serialize())
 
convert_onnx_to_trt_model('model.onnx', 'model.engine')



// 编写C++代码以使用TensorRT运行模型
// 需要TensorRT的C++ API支持
 
#include <iostream>
#include <cuda_runtime_api.h>
#include "NvInfer.h"
#include "NvOnnxParser.h"
 
int main() {
    // 创建TensorRT的引擎
    IBuilder* builder = createInferBuilder(gLogger);
    nvinfer1::INetworkDefinition* network = builder->createNetwork();
    nvonnxparser::IOnnxParser* parser = nvonnxparser::createParser(*network, gLogger);
 
    // 从engine文件加载TensorRT引擎
    ICudaEngine* engine = nullptr;
    engine = loadEngine("model.engine");
    if (!engine) {
        std::cout << "Unable to load engine" << std::endl;
        return -1;
    }
 
    // 运行模型的其余部分...
 
    // 清理资源
    parser->destroy();
    network->destroy();
    builder->destroy();
    engine->destroy();
 
    return 0;
}
最后修改于:2024年08月07日 21:00

评论已关闭

推荐阅读

DDPG 模型解析,附Pytorch完整代码
2024年11月24日
DQN 模型解析,附Pytorch完整代码
2024年11月24日
AIGC实战——Transformer模型
2024年12月01日
Socket TCP 和 UDP 编程基础(Python)
2024年11月30日
python , tcp , udp
如何使用 ChatGPT 进行学术润色?你需要这些指令
2024年12月01日
AI
最新 Python 调用 OpenAi 详细教程实现问答、图像合成、图像理解、语音合成、语音识别(详细教程)
2024年11月24日
ChatGPT 和 DALL·E 2 配合生成故事绘本
2024年12月01日
omegaconf,一个超强的 Python 库!
2024年11月24日
【视觉AIGC识别】误差特征、人脸伪造检测、其他类型假图检测
2024年12月01日
[超级详细]如何在深度学习训练模型过程中使用 GPU 加速
2024年11月29日
Python 物理引擎pymunk最完整教程
2024年11月27日
MediaPipe 人体姿态与手指关键点检测教程
2024年11月27日
深入了解 Taipy:Python 打造 Web 应用的全面教程
2024年11月26日
基于Transformer的时间序列预测模型
2024年11月25日
Python在金融大数据分析中的AI应用(股价分析、量化交易)实战
2024年11月25日
AIGC Gradio系列学习教程之Components
2024年12月01日
Python3 `asyncio` — 异步 I/O,事件循环和并发工具
2024年11月30日
llama-factory SFT系列教程:大模型在自定义数据集 LoRA 训练与部署
2024年12月01日
Python 多线程和多进程用法
2024年11月24日
Python socket详解,全网最全教程
2024年11月27日