yolov8实战第三天——yolov8TensorRT部署（python推理）（保姆教学）

在上一篇中我们使用自己的数据集训练了一个yolov8检测模型，best.py。

yolov8实战第一天——yolov8部署并训练自己的数据集（保姆式教程）-CSDN博客

yolov8实战第二天——yolov8训练结果分析（保姆式解读）-CSDN博客

接下要对best.py进行TensorRT优化并部署。

TensorRT是一种高性能深度学习推理优化器和运行时加速库，可以为深度学习应用提供低延迟、高吞吐率的部署推理。

TensorRT可用于对超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。

TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架，将TensorRT和NVIDIA的GPU结合起来，能在几乎所有的框架中进行快速和高效的部署推理。

一般的深度学习项目，训练时为了加快速度，会使用多GPU分布式训练。但在部署推理时，为了降低成本，往往使用单个GPU机器甚至嵌入式平台（比如 NVIDIA Jetson）进行部署，部署端也要有与训练时相同的深度学习环境，如caffe，TensorFlow等。

由于训练的网络模型可能会很大（比如，inception，resnet等），参数很多，而且部署端的机器性能存在差异，就会导致推理速度慢，延迟高。这对于那些高实时性的应用场合是致命的，比如自动驾驶要求实时目标检测，目标追踪等。

为了提高部署推理的速度，出现了很多模型优化的方法，如：模型压缩、剪枝、量化、知识蒸馏等，这些一般都是在训练阶段实现优化。

而TensorRT 则是对训练好的模型进行优化，通过优化网络计算图提高模型效率。

一、安装TensorRT  

Log in | NVIDIA Developer

下载TensorRT 。

我下载的是8.6里画黑线的那个。 

将 TensorRT-8.6.1.6\include中头文件 copy 到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\include
将TensorRT-8.6.1.6\lib 中所有lib文件 copy 到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\lib\x64
将TensorRT-8.6.1.6\lib 中所有dll文件copy 到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\bin

在python文件夹中找到适合自己的。

pip install tensorrt-8.6.1-cp310-none-win_amd64.whl

 至此TensorRT安装完成。

 二、pt转onnx：

GitHub - triple-Mu/YOLOv8-TensorRT: YOLOv8 using TensorRT accelerate !

参考着这个，下载，安装环境后。

安装onnx：

1. pip install onnx -i https://pypi.tuna.tsinghua.edu.cn/simple
2. pip install onnxsim -i https://pypi.tuna.tsinghua.edu.cn/simple
3. pip install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple

生成onnx： 

python export-det.py --weights yolov8n.pt --iou-thres 0.65 --conf-thres 0.25 --topk 100 --opset 11 --sim --input-shape 1 3 640 640 --device cuda:0

使用上篇文章中的老鼠模型做了测试： 

onnx的测试代码：

1. import onnxruntime as rt
2. import numpy as np
3. import cv2
4. import matplotlib.pyplot as plt
7. def nms(pred, conf_thres, iou_thres):
8. conf = pred[..., 4] > conf_thres
9. box = pred[conf == True]
10. cls_conf = box[..., 5:]
11. cls = []
12. for i in range(len(cls_conf)):
13. cls.append(int(np.argmax(cls_conf[i])))
14. total_cls = list(set(cls))
15. output_box = []
16. for i in range(len(total_cls)):
17. clss = total_cls[i]
18. cls_box = []
19. for j in range(len(cls)):
20. if cls[j] == clss:
21. box[j][5] = clss
22. cls_box.append(box[j][:6])
23. cls_box = np.array(cls_box)
24. box_conf = cls_box[..., 4]
25. box_conf_sort = np.argsort(box_conf)
26. max_conf_box = cls_box[box_conf_sort[len(box_conf) - 1]]
27. output_box.append(max_conf_box)
28. cls_box = np.delete(cls_box, 0, 0)
29. while len(cls_box) > 0:
30. max_conf_box = output_box[len(output_box) - 1]
31. del_index = []
32. for j in range(len(cls_box)):
33. current_box = cls_box[j]
34. interArea = getInter(max_conf_box, current_box)
35. iou = getIou(max_conf_box, current_box, interArea)
36. if iou > iou_thres:
37. del_index.append(j)
38. cls_box = np.delete(cls_box, del_index, 0)
39. if len(cls_box) > 0:
40. output_box.append(cls_box[0])
41. cls_box = np.delete(cls_box, 0, 0)
42. return output_box
45. def getIou(box1, box2, inter_area):
46. box1_area = box1[2] * box1[3]
47. box2_area = box2[2] * box2[3]
48. union = box1_area + box2_area - inter_area
49. iou = inter_area / union
50. return iou
53. def getInter(box1, box2):
54. box1_x1, box1_y1, box1_x2, box1_y2 = box1[0] - box1[2] / 2, box1[1] - box1[3] / 2, \
55. box1[0] + box1[2] / 2, box1[1] + box1[3] / 2
56. box2_x1, box2_y1, box2_x2, box2_y2 = box2[0] - box2[2] / 2, box2[1] - box1[3] / 2, \
57. box2[0] + box2[2] / 2, box2[1] + box2[3] / 2
58. if box1_x1 > box2_x2 or box1_x2 < box2_x1:
59. return 0
60. if box1_y1 > box2_y2 or box1_y2 < box2_y1:
61. return 0
62. x_list = [box1_x1, box1_x2, box2_x1, box2_x2]
63. x_list = np.sort(x_list)
64. x_inter = x_list[2] - x_list[1]
65. y_list = [box1_y1, box1_y2, box2_y1, box2_y2]
66. y_list = np.sort(y_list)
67. y_inter = y_list[2] - y_list[1]
68. inter = x_inter * y_inter
69. return inter
72. def draw(img, xscale, yscale, pred):
73. img_ = img.copy()
74. if len(pred):
75. for detect in pred:
76. detect = [int((detect[0] - detect[2] / 2) * xscale), int((detect[1] - detect[3] / 2) * yscale),
77. int((detect[0]+detect[2] / 2) * xscale), int((detect[1]+detect[3] / 2) * yscale)]
78. img_ = cv2.rectangle(img, (detect[0], detect[1]), (detect[2], detect[3]), (0, 255, 0), 1)
79. return img_
82. if __name__ == '__main__':
83. height, width = 640, 640
84. img0 = cv2.imread('mouse-4-6-0004.jpg')
85. x_scale = img0.shape[1] / width
86. y_scale = img0.shape[0] / height
87. img = img0 / 255.
88. img = cv2.resize(img, (width, height))
89. img = np.transpose(img, (2, 0, 1))
90. data = np.expand_dims(img, axis=0)
91. sess = rt.InferenceSession('best.onnx')
92. input_name = sess.get_inputs()[0].name
93. label_name = sess.get_outputs()[0].name
94. pred = sess.run([label_name], {input_name: data.astype(np.float32)})[0]
95. pred = np.squeeze(pred)
96. pred = np.transpose(pred, (1, 0))
97. pred_class = pred[..., 4:]
98. pred_conf = np.max(pred_class, axis=-1)
99. pred = np.insert(pred, 4, pred_conf, axis=-1)
100. result = nms(pred, 0.3, 0.45)
101. ret_img = draw(img0, x_scale, y_scale, result)
102. ret_img = ret_img[:, :, ::-1]
103. plt.imshow(ret_img)
104. plt.show()

三、TensorRT部署

导出engine模型：

python build.py --weights yolov8n.onnx --iou-thres 0.65 --conf-thres 0.25 --topk 100 --fp16 --device cuda:0

等待一会，engine成功导出。 

 

使用python脚本进行推理：

python infer-det.py --engine yolov8n.engine --imgs data --show --out-dir outputs --out-dir outputs --device cuda:0

infer-det.py: 

1. from models import TRTModule # isort:skip
2. import argparse
3. from pathlib import Path
5. import cv2
6. import torch
8. from config import CLASSES, COLORS
9. from models.torch_utils import det_postprocess
10. from models.utils import blob, letterbox, path_to_list
13. def main(args: argparse.Namespace) -> None:
14. device = torch.device(args.device)
15. Engine = TRTModule(args.engine, device)
16. H, W = Engine.inp_info[0].shape[-2:]
18. # set desired output names order
19. Engine.set_desired(['num_dets', 'bboxes', 'scores', 'labels'])
21. images = path_to_list(args.imgs)
22. save_path = Path(args.out_dir)
24. if not args.show and not save_path.exists():
25. save_path.mkdir(parents=True, exist_ok=True)
27. for image in images:
28. save_image = save_path / image.name
29. bgr = cv2.imread(str(image))
30. draw = bgr.copy()
31. bgr, ratio, dwdh = letterbox(bgr, (W, H))
32. rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
33. tensor = blob(rgb, return_seg=False)
34. dwdh = torch.asarray(dwdh * 2, dtype=torch.float32, device=device)
35. tensor = torch.asarray(tensor, device=device)
36. # inference
37. data = Engine(tensor)
39. bboxes, scores, labels = det_postprocess(data)
40. if bboxes.numel() == 0:
41. # if no bounding box
42. print(f'{image}: no object!')
43. continue
44. bboxes -= dwdh
45. bboxes /= ratio
47. for (bbox, score, label) in zip(bboxes, scores, labels):
48. bbox = bbox.round().int().tolist()
49. cls_id = int(label)
50. cls = CLASSES[cls_id]
51. color = COLORS[cls]
52. cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2)
53. cv2.putText(draw,
54. f'{cls}:{score:.3f}', (bbox[0], bbox[1] - 2),
55. cv2.FONT_HERSHEY_SIMPLEX,
56. 0.75, [225, 255, 255],
57. thickness=2)
58. if args.show:
59. cv2.imshow('result', draw)
60. cv2.waitKey(0)
61. else:
62. cv2.imwrite(str(save_image), draw)
65. def parse_args() -> argparse.Namespace:
66. parser = argparse.ArgumentParser()
67. parser.add_argument('--engine', type=str, help='Engine file')
68. parser.add_argument('--imgs', type=str, help='Images file')
69. parser.add_argument('--show',
70. action='store_true',
71. help='Show the detection results')
72. parser.add_argument('--out-dir',
73. type=str,
74. default='./output',
75. help='Path to output file')
76. parser.add_argument('--device',
77. type=str,
78. default='cuda:0',
79. help='TensorRT infer device')
80. args = parser.parse_args()
81. return args
84. if __name__ == '__main__':
85. args = parse_args()
86. main(args)