TensorRT8 meets Python(三) Onnx+TensorRT推理(案例基于yolov5 6.0)

 1.前言

        在前面两篇我们介绍了TensorRT的环境部署以及TensorRT的功能性介绍。

        在使用tensorRT的时候，最常见的方式就是各种训练框架都基于Onnx来做中间转换，通过Onnx来生成TensorRT engine,进而享受到TensorRT的推理速度。所以本文就以常见的检测模型yolov5 6.0 来进行TensorRT模型的部署。整体测试都是基于TensorRT python backend，在开始之间请大家思考如下问题？

        1.推理时所选的BatchSize 和 Fps的关系是怎样的？

        2.TensorRT推理一定比pytorch推理快吗，我们要无脑替换成TensorRT是正确的使用姿势吗？受到哪些因素的影响？对于使用python作为主要开发语言的同学来说，什么时候考虑使用tensorRT作为推理引擎，什么时候使用pytorch呢？

        带着这些问题，我们共同开始今天的学习。本文的实验环境是GPU V100S 32G，Intel(R) Xeon(R) Gold 6248R CPU @ 3.00GHz 。本文选用的方式是Explict batchSize，batchsize是从onnx模型继承过来的，关于Explicit 和 Implicit的方式，我们后续再进行介绍。

2.onnx转化nvidia engine的两种方法

2.1 onnx模型导出

        关于onnx模型导出的问题，我觉得还是另开篇章去讲述吧。export的代码yolov5官方已经提供了。大家可以参考一下这份代码：

yolov5官方onnx的转换代码

使用方法：

python export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite

重要参数：

• --weights : yolov5 pt模型的路径
• --batch-size: 推理的batchsize，这里使用静态的batchsize，必须在onnx内指定。
• --half：如果使用fp16精度进行推理，则需要指定这个参数。

我们需要测试不同batchsize下面的tensorRT模型，所以导出的时候指定不同的batchsize。（注意这是在Explicit batch模式下进行的）。另外我这边选择测试的是fp16，所以需要都加上--half参数。

2.2 基于Trtexec命令行工具

       最简单的方法就是使用trtexec作为onnx转换的工具。trtexec是Nvidia TensorRT自带的命令行工具。一方面可以帮助我们进行onnx或者其他格式的模型快速转换到TensorRT engine，包括指定batchsize、precision等，另一方面可以帮助我们测试模型的基准性能。

        因为我要测试不同batchsize下，tensorRT与pytorch之间的推理性能差距，所以在导出的时候需要使用对应的onnx模型。

trtexec --onnx=yolov5s.onnx  --saveEngine=yolov5s_engine_fp16_b64.trt  --inputIOFormats=fp16:chw --outputIOFormats=fp16:chw --explicitBatch --fp16

        这里因为是测试半精度，所以需要--inputIOFormats=fp16:chw --outputIOFormats=fp16:chw，这两个参数。我这里遇到一个坑就是，只指定--fp16，另外两个参数如果不指定的话，得到的模型推理速度会很慢。查阅文档是说默认是按照fp32的精度进行，即使单独指定了--fp16参数。

        我分别测试了4、8、16、32、64五组batchsize，得到了对应的5个engine模型。

2.3 基于tensorRT的python接口进行转换

另一种方式是基于tensorRT的python接口进行engine转换。

EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def ONNX_to_TRT(onnx_model_path=None, trt_engine_path=None, fp16_mode=False):
    """
    仅适用TensorRT V8版本
    生成cudaEngine，并保存引擎文件(仅支持固定输入尺度)
    Serialized engines不是跨平台和tensorRT版本的
    fp16_mode: True则fp16预测
    onnx_model_path: 将加载的onnx权重路径
    trt_engine_path: trt引擎文件保存路径
    """
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(EXPLICIT_BATCH)
    parser = trt.OnnxParser(network, TRT_LOGGER)
    config = builder.create_builder_config()
    config.max_workspace_size = GiB(1)
    if fp16_mode:
        config.set_flag(trt.BuilderFlag.FP16)
    with open(onnx_model_path, 'rb') as model:
        assert parser.parse(model.read())
        serialized_engine = builder.build_serialized_network(network, config)

    with open(trt_engine_path, 'wb') as f:
        f.write(serialized_engine)  # 序列化

    print('TensorRT file in ' + trt_engine_path)
    print('============ONNX->TensorRT SUCCESS============')

3. python调用tensorRT接口进行推理

class TrtModel():
    '''
    TensorRT infer
    '''

    def __init__(self, trt_path, device=1):
        self.cfx = cuda.Device(device).make_context()
        self.stream = cuda.Stream()
        TRT_LOGGER = trt.Logger(trt.Logger.INFO)
        # 启动一个tensorRT pyton runtime
        runtime = trt.Runtime(TRT_LOGGER)
        # 反序列化模型，使用runtime加载模型
        with open(trt_path, "rb") as f:
            self.engine = runtime.deserialize_cuda_engine(f.read())
        # 创建执行上下文
        self.context = self.engine.create_execution_context()
        # 模型输入的尺寸
        intype = self.engine.get_binding_dtype("input")
        insize = trt.volume(self.engine.get_binding_shape("input"))
        # fp16 占2个字节 fp32 占4个字节
        insize = insize * 2 if intype == DataType.HALF else insize * 4
        # 模型输出的尺寸
        otype = self.engine.get_binding_dtype("output")
        osize = trt.volume(self.engine.get_binding_shape("output"))
        osize = osize * 2 if otype == DataType.HALF else osize * 4
        otype = np.float16 if otype == DataType.HALF else np.float32
        # 分配输入输出的显存
        self.cuda_mem_input = cuda.mem_alloc(insize)
        self.cuda_mem_output = cuda.mem_alloc(osize)
        self.bindings = [int(self.cuda_mem_input), int(self.cuda_mem_output)]
        self.output = np.empty(self.engine.get_binding_shape("output"), dtype=otype)


    def __call__(self, img_np_nchw):
        '''
        TensorRT推理
        :param img_np_nchw: 输入图像
        '''
        self.cfx.push()
        #将数据从内存拷贝到显存中
        cuda.memcpy_htod_async(self.cuda_mem_input, img_np_nchw.ravel(), self.stream)
        #tensorRT异步并发推理
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        #将数据从显存传输到内存
        cuda.memcpy_dtoh_async(self.output, self.cuda_mem_output, self.stream)
        # 等待所有cuda核完成计算
        self.stream.synchronize()
        self.cfx.pop()
        return self.output

    def destroy(self):
        # Remove any context from the top of the context stack, deactivating it.
        self.cfx.pop()

        这是一个简化版本的tensorRT推理类。其原理其实就是提前开辟显存空间，然后把数据从内存拷贝到显存。在显卡上面利用cuda核进行并行计算，然后在从显存把结果数据拷贝到内存。

        另外关于pagelock memory（锁页内存）的使用，有些代码里会提前开辟锁页内存，我在实际测试的时候发现，并没有特别的速度提升。

4. 关于实验结论

        我在实际测试的时候发现，对于不同的batchsize，在tensorRT的python backend上面，不是完全都呈现出比pytorch推理更快的现象，而是呈现出有规律的趋势。在batchsize比较小的时候，tensorRT推理具备碾压的性能优势，但是随着batchsize的增大，这个gap在慢慢变小，在大于32、64等的情况下，反而没有pytorch直接推理快。

        在c++上面的测试，整体上tensorRT的推理都比pytorch快，但是同样这个gap在减少。

        有知道原因的朋友，求点拨。下面是我记录的实验记录。

         可以明显的看到这个趋势。所以在用python封装backend的时候，感觉tensorRT更适合小batchsize的使用场景，比如单张图片的某种推理服务。如果是视频推理的情景，可能不间断要一批批帧画面进行推理，对于这种情况，用pytorch的半精度直接推理可能效果反而更好，并不是无脑使用的。

 以上是以yolov5 6.0做的实验，来记录python backend的使用方法和记录测试结果。谢谢阅读。