TensorRT 模型推理优化实现

RedZero9

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