onnx 转 trt 并使用trt模型推理

最新推荐文章于 2024-04-11 13:43:05 发布

嘿，不许笑

最新推荐文章于 2024-04-11 13:43:05 发布

阅读量3.9k

点赞数

分类专栏： python深度学习文章标签： python 计算机视觉 opencv

本文链接：https://blog.csdn.net/qq_43994782/article/details/121536473

版权

本文记录了如何将ONNX模型转换为TensorRT（TRT）模型，并使用TRT模型进行推理。在转换过程中，提到了使用`trtexec`命令行工具以及TensorRT官方提供的Python API。在测试TRT模型时，遇到了一个因多余代码导致内存占用增加的有趣bug，原因是引入torch后在创建TRT运行时进行了不必要的内存预分配，注释掉这部分代码即可解决此问题。

摘要由CSDN通过智能技术生成

命令：

/usr/src/tensorrt/bin/trtexec --onnx=input.onnx --saveEngine=out.engine --fp16 --workspace=1024

参数含义：
/usr/src/tensorrt/bin/trtexec：你的tensorrt的地址
–onnx: 需要转换的onnx模型路径
–saveEngine: 保存的模型地址名称
–fp16: 保存出的模型的精度
–workspace: 运行空间大小（貌似没卵用）

参考nVidia tensorrt 官网

当然官网上也给了使用代码转换的案例包括c++和python
网址：https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#python_topics

但是
官方的案例并米有跑通，但是我从另一个博客里粘的代码调整了下，可以跑了，这里粘出来（不过还是推荐使用trtexec的方法。）

import tensorrt as trt


# TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
# builder = trt.Builder(TRT_LOGGER)

# EXPLICIT_BATCH = []
# network = builder.create_network(*EXPLICIT_BATCH)# = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
# config = builder.create_builder_config()
# parser = trt.OnnxParser(network, TRT_LOGGER)

# config.max_workspace_size = 1 << 30
# # Load the Onnx model and parse it in order to populate the TensorRT network.
# with open("det.onnx", 'rb') as model:
#     if not parser.parse(model.read()):
#         print ('ERROR: Failed to parse the ONNX file.')
#         for error in range(parser.num_errors):
#             print (parser.get_error(error))

# builder.build_engine(network, config)

import os
onnx_file_path = "det.onnx"
engine_file_path = 'det.engine'
def buildtrt():
    explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    # In TensorRT 7.0, the ONNX parser only supports full-dimensions mode, meaning that your network definition must be created with the explicitBatch flag set. For more information, see Working With Dynamic Shapes.
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
    with trt.Builder(TRT_LOGGER) as builder, \
        builder.create_network(explicit_batch) as network,  \
        trt.OnnxParser(network, TRT_LOGGER) as parser: