TensorRT 模型推理优化实现

快递包裹分割模块是单件分离系统中起决定性作用的一个环节，快递包裹分割算法的准确度和推理速度都影响着该系统的整体性能。想要将 TensorRT 优化之后的快递包裹分割算法部署到嵌入式开发板上需要经历三个阶段，分别为网络模型的训练阶段、网络模型格式的转换阶段和 TensorRT 推理优化器对模型结构的优化阶段。

1  网络模型存储格式转换 

      在使用 TensorRT 推理优化器对模型结构优化之前，需要先把网络模型结构和权重参数导入到推理优化器中。当前，TensorRT 包含 nvcaffeparse、nvuffparse和 nvonnxparse 三种文件解析器，可以支持对.caffe 文件、.uff 文件和.onnx 文件的解析。快递包裹分割算法是基于 PyTorch 深度学习框架训练的神经网络模型，因此，可以把训练好的网络模型结构和权重参数以 ONNX 的格式保存为.onnx 文件。 

     目前，PyTorch 深度学习框架可以使用 torch.onnx.export 函数把训练好的网络模型结构和权重参数保存为.onnx 文件。但是，在使用上述函数导出网络模型结构和权重参数之前，需要考虑 ONNX 是否支持网络结构中的所有运算操作以及网络模型中的后处理操作是否也要一起转换到.onnx 文件中。对于 ONNX 不支持的一些操作，通常有两种解决方法：第一种方法是使用其他近似操作来替换；第二种方法是把 PyTorch 定义的操作替换为自己实现的操作。对于使用自定义的操作来说，需要考虑 TesnorRT 框架是否支持和执行效率的问题，因此，在网络模型格式转换的过程中尽量使用近似的操作来替换。对于网络模型的后处理过程，通常都会存在着像非极大值抑制这样的操作，这种操作中存在着一些逻辑判断（比如 if 语句），ONNX 格式保存的网络模型文件对于这些逻辑运算是无能为力的。同时，对于后处理的实现方式，每个人操作可能都是不一样的，这对于使用TesnorRT 推理优化器优化网络模型结构增加了难度。总之，为了降低网络模型结构转换过程中可能存在的风险，后处理操作最好还是由开发者自己实现。 

2 TensorRT 优化模型结构

     使用 TensorRT 推理框架部署网络模型的过程中，必须要首先使用 TensorRT推理优化器对 ONNX 模型结构优化处理并生成 TensorRT 运行时引擎之后，才能对网络模型输入数据，然后执行推理操作获取最终的输出结果。但是，TensorRT推理优化器优化网络模型结构的过程非常耗时，所以，在实际的应用过程中都是先单独生成 TensorRT 推理引擎，然后再序列化为二进制文件保存在本地。应用时，只需要反序列化保存在本地的二进制文件，就可以快速地生成推理引擎，节约了大量的时间。 

TensorRT 推理引擎的构建过程如下图所示，具体步骤如下： 

1.  使用 nvinfer1::createInferBuilder 函数创建一个 nvinfer1::IBuilder 对象builder; 

2.  在 builder 对 象 上 调 用 createNetworkV2 函 数 创 建 一 个nvinfer1::INetworkDefinition 对象 network; 

3.  通过调用 nvonnxparse::createParser 函数并传入 network 参数创建一个nvonnxparser::IParser 对象 parser 来解析导入的 ONNX 网络模型文件； 

4.  在 builder 对 象 上 调 用 createBuilderConfig 函 数 创 建 一 个nvinfer1::IBuilderConfig 对象 config，通过这个对象可以设置一些在模型推理优化过程中使用到的参数，例如，网络模型的工作空间大小、每次网络模型推理时输入图片的数量和网络模型推理时使用的数据类型等等； 

5.  在 builder 对象上调用 builderEngineWithConfig 函数并传入 network 和config 参数，就可以对 ONNX 网络模型结构和权重参数进行优化推理和数据类型校准，最后生成 nvinfer1::ICudaEngine 对象 mEngine，也就是 TensorRT 推理优化器处理之后生成的推理引擎； 

6.  把生成的推理引擎 mEngine 序列化为二进制文件并存储在磁盘中，方便在后续模型部署的过程中直接调用。