※限制

• 如果使用4090顯卡去轉換, 最終在3060顯卡去使用會弱化效能.
• PyTorch的模型設定需要跟轉化後的模型相同設定.

方法1:

1. pt 轉 onnx

from ultralytics import YOLO
import os

os.chdir(os.path.dirname(os.path.abspath(__file__)))

'''
pip install onnx_graphsurgeon  # 用於處理 ONNX 模型的工具
pip install onnx  # ONNX 格式處理
pip install onnx-simplifier  # 用於簡化 ONNX 模型
pip install onnxruntime
pip install onnxslim
'''



weights_dir = './Weights'
# 載入 YOLOv8 模型，指定模型權重路徑
model_path = './Weights/yolov8n.pt'
model = YOLO(model_path)

model.export(format='onnx', imgsz=640)

2. onnx 轉 engine

import tensorrt as trt
import os
os.chdir(os.path.dirname(os.path.abspath(__file__)))

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
runtime = trt.Runtime(TRT_LOGGER)

# 讀取 ONNX 模型
onnx_file_path = './Weights/yolov8n.onnx'
with open(onnx_file_path, 'rb') as f:
    onnx_model = f.read()

# 轉換 ONNX 模型為 TensorRT 引擎
network_flags = 0
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(network_flags)
parser = trt.OnnxParser(network, TRT_LOGGER)

# 解析 ONNX 模型
if not parser.parse(onnx_model):
    print('ERROR: Failed to parse the ONNX model')
    for error in range(parser.num_errors):
        print(parser.get_error(error))
    exit(1)

# 生成 TensorRT 引擎
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
config.set_flag(trt.BuilderFlag.FP16)
engine = builder.build_serialized_network(network, config)


# 保存引擎
with open('./Weights/yolov8n.engine', 'wb') as f:
    f.write(engine)

方法2:
使用yolo 直接轉 engine

from ultralytics import YOLO
import os
'''
pip install onnxruntime-gpu
'''
os.chdir(os.path.dirname(os.path.abspath(__file__)))

model = YOLO("./Weights/yolov8n.pt")

# fp32
model.export(format="engine")

# fp16
model.export(format="engine", half=True)

# int8
model.export(
    format="engine",
    dynamic=True,  
    batch=8,  
    workspace=4,  
    int8=True,
    data="coco.yaml",
)