源码下载

https://gitee.com/yang_guo123/yolox_tensorrt
才疏学浅，若有不足之处请大家指出，喜欢的话可以点个star

一、整体实现方法

1.部署流程

在深度学习部署过程中，流程如下所示：

1、采集工程所需训练数据
2、通过pytorch、tensorflow等深度学习框架训练网络
3、将训练好的网络导出onnx格式文件
4、通过深度学习部署框架对onnx文件进行部署

2.TensorRT的安装

TensorRT是英伟达开发的一个深度学习推理框架，在N卡上具有较好的性能，在安装时，首先打开TensorRT主页（https://developer.nvidia.com/tensorrt），登录英伟达账号，并下载tensorrt8，在这里我使用的是8.2.4.2版本。

将下载的文件解压到自己的路径当中（注意：路径当中最好不带中文字符），接下来设置环境变量，将TensorRT主目录中的lib添加到环境变量当中，由于之后会使用trtexec工具，建议将bin文件夹也加入环境变量。

这时TensorRT的主体部分已经安装完成，接下来安装tensorrt在python中的库。
首先在anaconda当中创建一个tensorrt的环境，并使用pip安装pycuda库，用于调用GPU。

然后打开tensorrt主目录下的python文件夹，其中有python3.6-3.9的安装包，选择适合的安装包进行安装。

最后，在cmd当中打开python，并输入import tensorrt as trt，若没有报错，则说明安装成功。

3.onnx文件的导出

本文使用的yolox代码为https://github.com/bubbliiiing/yolox-pytorch, 因此接下来的模型的转换是以这个仓库作为教程，若使用官方代码，则需一定修改。
首先，将代码克隆到本地，并下载对应的权重文件。在predict.py文件中将mode设置为export_onnx，并运行代码导出onnx文件。

if __name__ == "__main__":
    yolo = YOLO(model_path=权重文件路径, classes_path=分类标签文件路径, phi='s')
    #----------------------------------------------------------------------------------------------------------#
    #   mode用于指定测试的模式：
    #   'predict'           表示单张图片预测，如果想对预测过程进行修改，如保存图片，截取对象等，可以先看下方详细的注释
    #   'video'             表示视频检测，可调用摄像头或者视频进行检测，详情查看下方注释。
    #   'fps'               表示测试fps，使用的图片是img里面的street.jpg，详情查看下方注释。
    #   'dir_predict'       表示遍历文件夹进行检测并保存。默认遍历img文件夹，保存img_out文件夹，详情查看下方注释。
    #   'heatmap'           表示进行预测结果的热力图可视化，详情查看下方注释。
    #   'export_onnx'       表示将模型导出为onnx，需要pytorch1.7.1以上。
    #----------------------------------------------------------------------------------------------------------#
    mode = "predict"

4、tensorrt引擎文件的生成

trtexec是tensorrt的一个自带的模型转换工具，用于转换模型、验证模型，在此只进行转换引擎文件的演示，若需要其他操作，可以输入trtexec -h命令，查看说明。
若你将tensorrt的bin目录添加入环境变量，则生成较为容易，在存放onnx文件的位置打开命令行窗口，输入trtexec --onnx=你的onnx文件路径 --saveEngine=生成文件的名称.engine 稍作等待，即可生成引擎文件。若你没有将bin目录加入环境变量，则首先需要切换至bin文件夹下进行操作。

trtexec --onnx=yolox_s.onnx --saveEngine=yolox.engine

注意：引擎文件的生成与运行trtexec的电脑硬件相关，生成的引擎文件只可以在生成的电脑中运行，即使有相同配置的电脑，也有可能运行不了别的电脑生成的引擎文件。

5、tensorrt部署网络

在部署网络时，我们需要使用到tensorrt中samples/python文件夹中的common.py，这个文件包装了部署网络所需的函数，其中以下是部署中需要使用的函数：

# 此函数是为引擎开辟缓存空间
# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # Allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        # Append the device buffer to device bindings.
        bindings.append(int(device_mem))
        # Append to the appropriate list.
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))
    return inputs, outputs, bindings, stream

# 此函数的功能是进行网络推理
# This function is generalized for multiple inputs/outputs for full dimension networks.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference_v2(context, bindings, inputs, outputs, stream):
    # Transfer input data to the GPU.
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # Run inference.
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # Transfer predictions back from the GPU.
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # Synchronize the stream
    stream.synchronize()
    # Return only the host outputs.
    return [out.host for out in outputs]

（1）初始化tensorrt

部署的第一步是初始化tensorrt的环境，包括创建日志对象、读取引擎文件以及开辟缓存空间，以下为初始化代码：

def __init__(self, engine_path, config, logger_severity, warmup_epoch=10):
    """
    YOLO类
    :param engine_path: engine文件路径
    :param config: config文件路径
    :param logger_severity: logger等级
    :param warmup_epoch: 预热轮次
    """
    self.logger = trt.Logger(logger_severity)    # 创建日志对象
    # 生成引擎以及开辟内存空间
    self.engine = self.get_engine(engine_path)
    self.context = self.engine.create_execution_context()
    self.inputs, self.outputs, self.binding, self.stream = common.allocate_buffers(self.engine)

    # 读取config文件
    with open(config, "r") as fp:
        self.config = yaml.safe_load(fp)

    self.num_classes = len(self.config["classes"])
    self.colors = get_n_hls_colors_v2(self.num_classes)    # 获取检测框颜色

    # 网络预热
    self.warmup(warmup_epoch)

def get_engine(self, engine_path):
    # If a serialized engine exists, use it instead of building an engine.
    print("Reading engine from file {}".format(engine_path))
    with open(engine_path, "rb") as f, trt.Runtime(self.logger) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

（2）数据预处理

第二部是读取图像并进行预处理，在图像输入网络前，需要对其进行预处理，将图像resize成网络所需的尺寸，并进行归一化、标准化，这部分在yolo.py文件中的transform函数中：

def transform(self, img):
    """
    图像预处理
    :param img: 输入图像
    :return: 预处理后的图像
    """
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = letterbox_image(img, self.config["input_shape"][2:]).astype(np.float32) / 255.    # 给图像增加灰条并resize

    # 标准化图像
    img -= np.array(self.config["mean"])
    img /= np.array(self.config["std"])

    # shape: 640x640x3 -> 1x3x640x640
    img = np.transpose(img, [2, 0, 1])
    img = np.expand_dims(img, axis=0)
    return img

（3）网络推理

第三步则是将图像送入网络进行推理，这一步需要进行三个操作：将图像送入指定内存，网络推理以及从网络推理结果当中获取特征层：

np.copyto(self.inputs[0].host, img.reshape(-1))    # 将图像拷贝到分配的内存当中
result = common.do_inference_v2(self.context, self.binding, self.inputs, self.outputs, self.stream)  # 网络推理
# 获取输出特征层
for i in range(3):
    result[i] = np.reshape(result[i], newshape=[1, 5 + self.num_classes] + self.config["stage"][i])

（4）网络后处理

由于使用的库在导出时没有进行后处理，所以在网络推理后，还需进行后处理，将检测框从中提取出来，这一步由于代码较长，可以参考代码以及bubbliiiig大佬的yolox文章。

二、代码库的使用

首先，使用git命令将代码从gitee上拉取到本地（由于GitHub网速缓慢，因此使用gitee）

git clone https://gitee.com/yang_guo123/yolox_tensorrt.git

我在百度网盘当中存放了通过coco数据集训练的onnx文件，若是只想进行测试则可以使用
链接：https://pan.baidu.com/s/1HuDrb1xtGKF14qM5nOZekQ 提取码：yolo

然后通过上文第四步在自己的电脑当中生成tensorrt所需的引擎文件，并放入engine文件夹当中。
下一步，则需修改config文件夹中的yolox.yaml，或者也可根据格式重新创建（若是使用百度网盘中的onnx文件此步可省略）

input_shape: [1, 3, 640, 640]      # 输入网络的尺寸

mean: [0.485, 0.456, 0.406]      # 均值
std: [0.229, 0.224, 0.225]      # 方差

stage: [[80, 80], [40, 40], [20, 20]]      # 网络生成的三个特征层的尺寸

conf_thres: 0.3      # 置信度阈值
iou_thres: 0.5      # iou阈值

# 类别名称
classes: [ 'person', 'bicycle', 'car', 'motorbike', 'aeroplane', 'bus', 'train',
           'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter',
           'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
           'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase',
           'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
           'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife',
           'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog',
           'pizza', 'donut', 'cake', 'chair', 'sofa', 'pottedplant', 'bed', 'diningtable', 'toilet',
           'tvmonitor', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
           'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
           'hair drier', 'toothbrush' ]

最后，运行inference.py文件，进行网络的推理。

总结

这篇文章主要讲述如何通过tensorrt进行网络推理，以及我的工程的使用，下一节将会讲述如何将yolox部署到jetson开发板中

reference

1、https://blog.csdn.net/weixin_44791964/article/details/120476949?spm=1001.2014.3001.5502
2、https://github.com/NVIDIA/trt-samples-for-hackathon-cn/tree/master/cookbook