C++ Tensorrt部署yolov8分割模型流程 自用留痕
·
环境配置: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博客感谢分享!
更多推荐
所有评论(0)