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 trt
 trt.__version__

转换

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

from ultralytics import YOLO

model = 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 trt


def build_engine(onnx_file_path, engine_file_path):
    # 创建TensorRT logger对象
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

    # 创建builder和network
    builder = 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) 并且最大批量大小是 8
    input_name = network.get_input(0).name
    profile.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)

参考：文章