Yolov8将.pt文件转换为tensorRt的.trt文件（模型部署）

本文主要是介绍Yolov8将.pt文件转换为tensorRt的.trt文件（模型部署），希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

我的环境

确保自己已经有cuda和cudnn的环境基础上进行。

cuda:11.7
cudnn:适合cuda的版本
Anaconda3 [python 3.10]
TensorRt-8.6.1

安装TensorRt环境

查看自己的cuda环境，去官网下载适合的win版本。

官网地址

下载后解压，将解压后lib目录添加到环境变量。

找到自己的cuda安装文件目录，将解压后lib文件夹下的lib文件，拷贝到cuda安装目录的x64文件（我这里路径是/cuda/lib/x64）中。

找到自己的cuda安装文件目录，将解压后lib文件夹下的dll文件，拷贝到cuda安装目录下bin的文件中（我这里路径是/cuda/bin）中。

查看python版本并安装对应的tensorrt。tensorrt文件在解压后的python目录。

选择适合的版本。

我的安装命令

pip install "D:\Programming software\TensorRt\TensorRT-8.6.1.6.Windows10.x86_64.cuda-11.8\TensorRT-8.6.1.6\python\tensorrt-8.6.1-cp310-none-win_amd64.whl"

测试

 # 测试import tensorrt as trttrt.__version__

转换

将Yolov8中.pt文件转换为onnx格式

from ultralytics import YOLOmodel = YOLO("best.pt")success = model.export(format="onnx", half=False, dynamic=True, opset=17)print("demo")

将onnx格式转换为tensorRt中的.trt格式。注意，我这里版本为tensorrt-8.6.1,其他版本可能Api 不同。

import tensorrt as trtdef build_engine(onnx_file_path, engine_file_path):# 创建TensorRT logger对象TRT_LOGGER = trt.Logger(trt.Logger.WARNING)# 创建builder和networkbuilder = trt.Builder(TRT_LOGGER)config = builder.create_builder_config()profile = builder.create_optimization_profile()flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)network = builder.create_network(flag)# 创建ONNX解析器parser = trt.OnnxParser(network, TRT_LOGGER)# 解析ONNX模型with open(onnx_file_path, 'rb') as onnx_file:if not parser.parse(onnx_file.read()):print('ERROR: Failed to parse the ONNX file.')for error in range(parser.num_errors):print(parser.get_error(error))return None# 设置最大工作空间config.max_workspace_size = 1 << 30  # 1GB# 选择模型的输入和输出格式，以及相关优化# 如果你知道模型运行的最大批量大小，你可以在这里设置# 例如: 如果你的模型输入尺寸是 CHW 的 (3, 224, 224) 并且最大批量大小是 8input_name = network.get_input(0).nameprofile.set_shape(input_name, min=(1, 3, 224, 224), opt=(4, 3, 224, 224), max=(8, 3, 224, 224))config.add_optimization_profile(profile)# 根据配置构建引擎engine = builder.build_engine(network, config)# 序列化引擎并保存到文件with open(engine_file_path, 'wb') as engine_file:engine_file.write(engine.serialize())return engine# ONNX模型和TensorRT引擎文件的路径
onnx_model_path = 'best.onnx'
tensorrt_engine_path = 'model.trt'# 构建并保存TensorRT引擎
engine = build_engine(onnx_model_path, tensorrt_engine_path)
if engine:print("ONNX model has been successfully converted to TensorRT engine and saved to", tensorrt_engine_path)