YOLOv5-6.1从训练到部署（四）：模型在GPU上部署_yolov5 gpu

上一篇文章讲到了如何使用OpenCV DNN和OpenVINO将模型部署到CPU上，但实际应用场景中，可能部署设备上存在GPU，此时我们可以使用ONNXRUNTIME或者TensorRT进行部署。

1 yolov5官方程序的GPU推理时间

如何证明我们自己写的部署程序比官方推理程序快？总不可能我们写的程序放GPU上，官方放CPU上吧，这样太不公平了，要对比就放同一设备上。本机使用RTX 3060显卡，我们就来测量一下在GPU上的时间。

在命令行输入：

python detect.py --source uav_bird_training/data/images/train/20220318_01.jpg --weights runs/train/exp2/weights/best.pt --data uav_bird_training/dataset.yaml --device 0
1

输出：

可以看到，使用GPU的情况下，预处理+正向推理+后处理，三步只需要20ms，已经是非常快了。

2 ONNXRUNTIME部署

ONNX Runtime（ONNX Runtime，简称ORT）是微软推出的用于深度学习模型推理（inference）的高性能开源推理引擎，本节将介绍如何通过ONNXRUNTIME，将我们训练好的yolov5模型部署到GPU中。

在yolov5-6.1目录下创建一个名为inference_openvino.py的文件，文件内部创建一个名为Inference_Onnxruntime的类，该类的结构与上一篇文章中的Inference_Opencv基本一样，改变的只有__init__和pred_img两个类内函数，并且也只是稍微改动了一点而已。这两个函数的代码如下：

import os
import cv2
import time
import yaml
import numpy as np
import onnxruntime


class Inference_Onnxruntime():
    # 全局设置（也可以在__init__中将它们设置成实例属性）
    INPUT_WIDTH = 640
    INPUT_HEIGHT = 640

    def __init__(self, onnx_path, yaml_path, score_threshold=0.25, nms_threshold=0.45, out_dir='out'):
        """
        初始化方法
        Args:
            onnx_path: onnx文件路径
            yaml_path: 数据集配置文件路径，这里主要是通过它来获取数据集有哪些类别
            score_threshold: 置信度得分阈值
            nms_threshold: NMS时的IOU阈值
            out_dir: 检测结果保存目录，暂时只能保存图像，摄像头/视频后续可以加
        """
        # 获取类列表
        with open(yaml_path, "r", errors='ignore') as f:
            self.class_list = yaml.safe_load(f)['names']

        # 创建推理会话
        self.session = onnxruntime.InferenceSession(onnx_path,
                                                    providers=['CUDAExecutionProvider'])

        # 绘制预测框、文字所用的颜色
        self.colors = [(255, 255, 0), (0, 255, 0), (0, 255, 255), (255, 0, 0)]

        # 预测框过滤相关的阈值设置
        self.score_threshold = score_threshold
        self.nms_threshold = nms_threshold

        # 检测结果保存目录
        self.out_dir = out_dir
        if not os.path.exists(self.out_dir):
            os.makedirs(self.out_dir)

    def pred_img(self, img_path):
        """
        预测图像
        Args:
            img_path: 图像路径
        """
        start = time.time()

        # 读取图像并预处理
        image = cv2.imread(img_path)
        time1 = time.time()
        print('read:', time1 - start)

        inputImage, factor, (dh, dw) = self.preprocess(image,
                                                       (Inference_Onnxruntime.INPUT_HEIGHT, Inference_Onnxruntime.INPUT_WIDTH))
        time2 = time.time()
        print('preprocess:', time2 - time1)

        # 模型正向推理
        ort_inputs = {self.session.get_inputs()[0].name: inputImage}
        outs = self.session.run(None, ort_inputs)[0]

        time3 = time.time()
        print('refer:', time3 - time2)

        # 解析推理结果（后处理）
        class_ids, scores, boxes = self.wrap_detection2(outs[0])	# self.wrap_detection2内部使用numpy广播机制
        time4 = time.time()
        print('wrap_detection:', time4 - time3)

        # 绘图
        image = self.draw_boxes(image, factor, (dh, dw), class_ids, scores, boxes)
        time5 = time.time()
        print('draw boxes:', time5 - time4)

        # 计算fps
        end = time.time()
        inf_end = end - start
        fps = 1 / inf_end
        fps_label = "FPS: %.2f" % fps
        cv2.putText(image, fps_label, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

        time6 = time.time()
        print('compute fps:', time6 - time5)

        # 保存
        basename = os.path.basename(img_path)
        cv2.imwrite(os.path.join(self.out_dir, basename), image)

        time7 = time.time()
        print('save:', time7 - time6)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

其他关于前处理（preprocess）、后处理（wrap_detection2）、画框（draw_boxes）的函数，与Inference_Opencv类完全一致，这里不再赘述。

测试程序如下：

if __name__ == '__main__':
    onnx_path = "runs/train/exp2/weights/best.onnx"
    yaml_path = "uav_bird_training/dataset.yaml"
    inference_Model = Inference_Onnxruntime(onnx_path, yaml_path)

    # 对保存在磁盘上的图片进行推理
    start = time.time()
    inference_Model.pred_img('uav_bird_training/data/images/train/20220318_01.jpg')
    print(time.time() - start)
1
2
3
4
5
6
7
8
9

输出：

read: 0.002992391586303711
preprocess: 0.00997304916381836
refer: 1.2430415153503418
wrap_detection: 0.000997781753540039
draw boxes: 0.27842187881469727
compute fps: 0.0
save: 0.00498652458190918
1.5404131412506104
1
2
3
4
5
6
7
8

这里有两个蹊跷的地方，首先，这里推理时间（refer）比之前所有的方案都慢，按理说都用GPU了，不应该慢才对；此外，画框的程序与之前一模一样但这里画框的时间却远远少于之前。

这是由于ONNXRUNTIME需要在第一次正向传播时建图，因此refer占用的时间很长。为了科学地统计时间，我们这里用第二次推理的时间来评估ONNXRUNTIME的推理速度，相关的测试代码如下：

if __name__ == '__main__':
    onnx_path = "runs/train/exp2/weights/best.onnx"
    yaml_path = "uav_bird_training/dataset.yaml"
    inference_Model = Inference_Onnxruntime(onnx_path, yaml_path)

    # 对保存在磁盘上的图片进行推理
    start = time.time()
    inference_Model.pred_img('uav_bird_training/data/images/train/20220318_01.jpg')
    print(time.time() - start)
    print('--------------------------------')

	# 第二次推理
	start = time.time()
    inference_Model.pred_img('uav_bird_training/data/images/train/20220318_01.jpg')
    print(time.time() - start)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

输出：

read: 0.002992391586303711
preprocess: 0.009973287582397461
refer: 1.261014699935913
wrap_detection: 0.0
draw boxes: 0.26511549949645996
compute fps: 0.0
save: 0.002991914749145508
1.5430963039398193
--------------------------------
read: 0.001983642578125
preprocess: 0.00498652458190918
refer: 0.008976221084594727
wrap_detection: 0.0
draw boxes: 0.0009975433349609375
compute fps: 0.0
save: 0.003988027572631836
0.02093195915222168
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

好的，从上面的结果来看，第二次检测图像时，refer所花的时间明显减少，我们因此看到了onnxruntime的性能。预处理、模型前向传播、后处理三步合计耗时为14ms，已经低于官方程序的GPU推理用时。

当然，我们同样可以用第二次推理的时间来衡量OpenCV DNN和OpenVINO这两个框架的推理用时，最后的结果和第一次推理的用时差不多，感兴趣的小伙伴自己可以尝试。

3 TensorRT部署

我们刚刚通过ONNXRUNTIME实现了模型的GPU部署模型，并且做到了比官方程序更快的速度，但我们希望推理速度能再快一点，此时可以考虑使用TensorRT。

TensorRT是nvidia家的一款高性能深度学习推理SDK。此SDK包含深度学习推理优化器和运行环境，可为深度学习推理应用提供低延迟和高吞吐量，在推理过程中，基于TensorRT的应用程序比仅仅使用CPU作为平台的应用程序要快40倍。

前面介绍的OpenCV DNN、OpenVINO和ONNXRUNTIME，它们既可以实现模型的CPU部署，也可以实现GPU，TensorRT与这些框架不同的是，它只能GPU部署，毕竟英伟达公司的主业不是CPU。

篇幅有限，这里只介绍Windows上的安装，如果用的Linux（比如用的事前面介绍过的AutoDL、AutoLn等），则可以跳过这一节，直接看这篇文章。

（1）TensorRT的下载

首先要查看CUDA版本，在命令行中输入如下命令：

nvcc --version
1

假如我们还想看自己装的cuDNN什么版本，可以先查看CUDA的安装路径：

where nvcc
1

可以看到本机的CUDA版本为11.1，安装目录为：C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.1

我们进到这个目录里，打开名为include的文件夹，然后在下面找到名为cudnn.h的文件：

用notepad++或者记事本打开，滑到最后，可以看到几行C++的宏定义：

从上面画框的信息来看，本机的cuDNN版本为8.2.1。

接下来是下载TensorRT，进入TensorRT的下载页面

在弹出的窗口中，输入自己的邮箱：

在TensorRT的列表中，我们选择TensorRT8

在需要同意的地方打钩

选择TensorRT 8.6 GA，如果CUDA是10.X，那么必须选8.5或者8.5以下的TensorRT，因为自8.6开始已经不再支持CUDA10.X了。

根据自己的平台选择程序包（我当前的电脑是Windows10，而且CUDA版本是11.0，因此选择红色画框部分）：

（2）TensorRT的配置（python）

下载之后解压，然后在解压后的目录下找到python：

进去之后根据本地的python环境找到对应的whl文件，这里我们选择完整的安装包（dispatch和lean都是不完整的）：

随后在这个目录下打开终端、激活环境、安装whl文件：

pip install tensorrt-8.6.1-cp36-none-win_amd64.whl
1

接下来安装安装onnx python sdk支持，相关的whl文件在解压文件的onnx_graphsuigeon中

在命令行中返回到上一级目录，随后进入到onnx_graphsurgeon中，去安装onnx python sdk支持

cd ..
cd onnx_graphsurgeon
pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl
1
2
3

接着安装uff的whl文件，它在TensorRT-8.6.1.6/uff下面

然后将TensorRT的lib目录配置到系统的环境变量中

然后将启动解释器测试版本

>>> import tensorrt
>>> tensorrt.__version__
'8.6.1'
>>>
1
2
3
4

至此，TensorRT的python sdk安装完成。

关于TensorRT的快速入门，可以看这篇文章，建议和GitHub上的TensorRT配合起来看：

在IntroNotebooks下面，有jupyter notebook的教程。

（3）TensorRT推理文件的导出

我们先要导出推理文件：

python export.py --weights runs/train/exp2/weights/best.pt --include engine --device 0
1

这里的--device 0是因为我这里只有一个GPU，如果有多张卡，可以换成其他数字，但不能省略--device这个参数，否则会默认为cpu，因为TensorRT只能面向GPU，因此省略会报错。

这个导出时间稍微长了一点，导出后，终端显示如下：

官方推理程序也可以使用刚刚导出的engine文件进行推理：

python detect.py --source uav_bird_training/data/images/train/20220318_01.jpg --weights runs/train/exp2/weights/best.engine --data uav_bird_training/dataset.yaml --device 0
1

结果：

使用TensorRT，推理时间大幅下降，预处理+前向传播+后处理只需要10ms。需要注意的事，这里生成的engine文件是和硬件相关的，不同型号的显卡不能通用这个engine。

当然，官方推理程序是很难部署的，有很多依赖的类别和库，我们要自己写一段部署程序。

（4）TensorRT的Python部署

inference_tensorrt.py的文件，里面新建一个名为Inference_TensorRT的类，该类的结构与Inference_Opencv基本一样，改变的只有__init__和pred_img两个类内函数，改动的代码也只是针对TensorRT框架的设置和推理的相关过程。这两个函数的代码如下：

import os
import cv2
import time
import yaml
import torch
import numpy as np
import tensorrt as trt
from collections import OrderedDict, namedtuple


class Inference_TensorRT():
    # 全局设置（也可以在__init__中将它们设置成实例属性）
    INPUT_WIDTH = 640
    INPUT_HEIGHT = 640

    def __init__(self, engine_path, device, yaml_path, score_threshold=0.25, nms_threshold=0.45, out_dir='out'):
        """
        初始化方法
        Args:
            engine_path: TensorRT引擎文件
            device: 推理设备
            yaml_path: 数据集配置文件路径，这里主要是通过它来获取数据集有哪些类别
            score_threshold: 置信度得分阈值
            nms_threshold: NMS时的IOU阈值
            out_dir: 检测结果保存目录，暂时只能保存图像，摄像头/视频后续可以加
        """
        # 获取类列表
        with open(yaml_path, "r", errors='ignore') as f:
            self.class_list = yaml.safe_load(f)['names']

        # 推理引擎的相关配置
        Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
        logger = trt.Logger(trt.Logger.INFO)
        with open(engine_path, 'rb') as f, trt.Runtime(logger) as runtime:
            self.engine = runtime.deserialize_cuda_engine(f.read())      # 注意这里的engine必须做成属性，
            # 即必须是self.engine，而不能是engine，虽然在self.pred_img中并没有直接使用model，但间接使用了
            # 如果这里不做成类的属性，那么在初始化方法结束后，engine将被释放，使得推理报错
        self.bindings = OrderedDict()
        for index in range(self.engine.num_bindings):
            name = self.engine.get_binding_name(index)
            dtype = trt.nptype(self.engine.get_binding_dtype(index))
            shape = self.engine.get_binding_shape(index)
            data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(device)
            self.bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
        self.binding_addrs = OrderedDict((n, d.ptr) for n, d in self.bindings.items())
        self.context = self.engine.create_execution_context()

        # 绘制预测框、文字所用的颜色
        self.colors = [(255, 255, 0), (0, 255, 0), (0, 255, 255), (255, 0, 0)]

        # 预测框过滤相关的阈值设置
        self.score_threshold = score_threshold
        self.nms_threshold = nms_threshold

        # 检测结果保存目录
        self.out_dir = out_dir
        if not os.path.exists(self.out_dir):
            os.makedirs(self.out_dir)

    def pred_img(self, img_path):
        """
        预测图像
        Args:
            img_path: 图像路径
        """
        start = time.time()

        # 读取图像
        image = cv2.imread(img_path)
        time1 = time.time()
        print('read:', time1 - start)

        # 预处理
        inputImage, factor, (dh, dw) = self.preprocess(image,
                                                       (Inference_TensorRT.INPUT_HEIGHT, Inference_TensorRT.INPUT_WIDTH))
        time2 = time.time()
        print('preprocess:', time2 - time1)

        # TensorRT推理
        x_input = torch.from_numpy(inputImage).to(device)
        self.binding_addrs['images'] = int(x_input.data_ptr())
        self.context.execute_v2(list(self.binding_addrs.values()))
        outs = self.bindings['output'].data.cpu().numpy()

        time3 = time.time()
        print('refer:', time3 - time2)

        # 解析推理结果（后处理）
        class_ids, scores, boxes = self.wrap_detection2(outs[0])	# self.wrap_detection2内部使用numpy广播机制
        time4 = time.time()
        print('wrap_detection:', time4 - time3)

        # 绘图
        image = self.draw_boxes(image, factor, (dh, dw), class_ids, scores, boxes)
        time5 = time.time()
        print('draw boxes:', time5 - time4)

        # 计算fps
        end = time.time()
        inf_end = end - start
        fps = 1 / inf_end
        fps_label = "FPS: %.2f" % fps
        cv2.putText(image, fps_label, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

        time6 = time.time()
        print('compute fps:', time6 - time5)

        # 保存
        basename = os.path.basename(img_path)
        cv2.imwrite(os.path.join(self.out_dir, basename), image)

        time7 = time.time()
        print('save:', time7 - time6)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

其他关于前处理（preprocess）、后处理（wrap_detection2）、画框（draw_boxes）的函数，与Inference_Opencv类完全一致。

测试程序如下：

if __name__ == '__main__':
    engine_path = "runs/train/exp2/weights/best.engine"
    yaml_path = "uav_bird_training/dataset.yaml"
    device = 'cuda:0'
    inference_Model = Inference_TensorRT(engine_path, device, yaml_path)

    # 对保存在磁盘上的图片进行推理
    start = time.time()
    inference_Model.pred_img('uav_bird_training/data/images/train/20220318_01.jpg')
    print(time.time() - start)
    print('---------------------------------')

    # 第二次推理
    start = time.time()
    inference_Model.pred_img('uav_bird_training/data/images/train/20220318_01.jpg')
    print(time.time() - start)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

输出：

[12/30/2023-02:29:32] [TRT] [I] Loaded engine size: 36 MiB
[12/30/2023-02:29:32] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in engine deserialization: CPU +0, GPU +33, now: CPU 0, GPU 33 (MiB)
[12/30/2023-02:29:33] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +33, now: CPU 0, GPU 66 (MiB)
[12/30/2023-02:29:33] [TRT] [W] CUDA lazy loading is not enabled. Enabling it can significantly reduce device memory usage and speed up TensorRT initialization. See "Lazy Loading" section of CUDA documentation https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#lazy-loading
read: 0.002992391586303711
preprocess: 0.0109710693359375
refer: 0.00698089599609375
wrap_detection: 0.0
draw boxes: 0.2449805736541748
compute fps: 0.0
save: 0.003989219665527344
0.2699141502380371
---------------------------------
read: 0.001994609832763672
preprocess: 0.004986763000488281
refer: 0.005983591079711914
wrap_detection: 0.0009980201721191406
draw boxes: 0.0
compute fps: 0.0
save: 0.003988981246948242
0.01795196533203125
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

我们看第二次推理的结果，预处理、模型前向传播、后处理三步合计耗时为不到11ms，其中模型前向传播耗时为5.98ms，相对于ONNXRUNTIME的8.98ms也是有所提升的。

（5）TensorRT的C++部署程序（暂时不需要掌握）

大部分公司都要求图像算法工程师会C++，因此掌握C++编程，已经成为了算法工程师的一项基本能力。实际应用场景中，需要考虑模型的性能和效率，比如运行速度、内存占用、功耗等，此时Python很难满足要求，所以在深度学习模型部署时，一般是使用C++语言，对于OpenVINO、ONNXRUNTIME部署模型，也普遍是使用C++语言。不过，由于本课程的学员普遍缺乏C++基础，因此C++部署部分暂时不要求掌握，这里就不展开介绍。

VS2017配置TensorRT，可以看B站的这个视频。

4 关于推理框架的总结

至此，我们已经学习了如何通过OpenCV DNN、OpenVINO、ONNXRUNTIME和TensorRT部署yolov5模型，这几个框架的使用流程大同小异，也各有优缺点，具体如何选择，可以参考下面的几条经验：
（1）如果模型需要部署在CPU或者英特尔的产品上，则优先选择OpenVINO；
（2）如果模型需要部署在GPU或者英伟达的产品上，则优先选择TensorRT；
（3）如果模型比较新，并且其中使用了比较新的算子，那么先尝试ONNXRUNTIME，因为ONNXRUNTIME兼容性最好，OpenVINO和TensorRT对最新算子的支持，可能存在滞后性，随后再根据硬件平台选择OpenVINO或TensorRT。

5 总结

本文介绍了如何使用ONNXRUNTIME和TensorRT，将模型部署到GPU上，本文的重点为：（1）通过ONNXRUNTIME进行深度学习模型GPU推理；（2）如何配置TensorRT，并导出TensorRT的推理文件（engine文件）；（3）通过TensorRT进行深度学习模型GPU推理；（4）根据模型的部署硬件，合理选择推理框架。考虑到很多学员都没有C++基础，因此使用C++部署模型，暂时不要求掌握。

本系列文章——YOLOv5-6.1从训练到部署，至此全部结束，这四篇文章前前后后花了我一个多月的时间，在写教程的时候，自己也查了不少资料，也学了很多新的工具，在这过程中，我自身的水平也获得了相应的提高。