环境配置:Cuda11.5、tensorrt8.4.1.5(高版本算子改变较多 稳定使用就用这版)、pytorch1.11.0+cu115

模型转换 pt->onnx->engine

用底部链接内export-seg.py 如果用yolo自带export会导致模型输出不匹配。pt->onnx:

import argparse
from io import BytesIO

import onnx
import torch
from ultralytics import YOLO

from models.common import optim

try:
    import onnxsim
except ImportError:
    onnxsim = None


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-w',
                        '--weights',
                        type=str,
                        default=r"D:\PythonProjects\YOLOv8-TensorRT-main\YOLOv8-TensorRT-main\model\1.6\best2.pt",
                        help='PyTorch yolov8 weights')
    parser.add_argument('--opset',
                        type=int,
                        default=13,
                        help='ONNX opset version')
    parser.add_argument('--sim',
                        action='store_true',
                        help='simplify onnx model')
    parser.add_argument('--input-shape',
                        nargs='+',
                        type=int,
                        default=[1, 3, 640, 640],
                        help='Model input shape only for api builder')
    parser.add_argument('--device',
                        type=str,
                        default='cpu',
                        help='Export ONNX device')
    args = parser.parse_args()
    assert len(args.input_shape) == 4
    return args


def main(args):
    YOLOv8 = YOLO(args.weights)
    model = YOLOv8.model.fuse().eval()
    for m in model.modules():
        optim(m)
        m.to(args.device)
    model.to(args.device)
    fake_input = torch.randn(args.input_shape).to(args.device)
    for _ in range(2):
        model(fake_input)
    save_path = args.weights.replace('.pt', '.onnx')
    with BytesIO() as f:
        torch.onnx.export(model,
                          fake_input,
                          f,
                          opset_version=args.opset,
                          input_names=['images'],
                          output_names=['outputs', 'proto'])
        f.seek(0)
        onnx_model = onnx.load(f)
    onnx.checker.check_model(onnx_model)
    if args.sim:
        try:
            onnx_model, check = onnxsim.simplify(onnx_model)
            assert check, 'assert check failed'
        except Exception as e:
            print(f'Simplifier failure: {e}')
    onnx.save(onnx_model, save_path)
    print(f'ONNX export success, saved as {save_path}')


if __name__ == '__main__':
    main(parse_args())

转换失败可能情况:找不到tensorrt的dll,检查tensorrt是否添加到环境,如果还是不行直接将缺少的dll拉到报错提示的路径下。

onnx->engine:

下载tensorrt后,在该路径下进入cmd

输入trtexec.exe --onnx=best.onnx --saveEngine=best.engine --fp16

回车后开始转换,ps:转换失败情况可能是pytorch或cuda的环境没配置好或者与tensorrt不匹配;也可能是onnx格式错误,可尝试重新生成或换上面代码生成。还有可能内存不够中断,关闭其他内存占用程序即可。

生成完成后,使用https://github.com/triple-Mu/YOLOv8-TensorRT内cmake生成部署代码,配置环境,修改路径调用即可。

参考:YOLOv8-TensorRT项目实战:基于TensorRT的YOLOv8实例分割模型部署指南-CSDN博客感谢分享!

Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐