【python】tensorrt8版本下的onnx转tensorrt engine

最新推荐文章于 2025-03-20 14:33:10 发布

无情的AI鸽子

最新推荐文章于 2025-03-20 14:33:10 发布

阅读量6.5k

点赞数 9

文章标签： python 深度学习

本文链接：https://blog.csdn.net/weixin_42492254/article/details/125319112

版权

博客介绍了如何在Python环境中使用TensorRT 8将ONNX模型转换为可执行的TensorRT引擎。作者在过程中遇到了CUDA和TRT版本冲突的问题，解决后尝试通过python onnx2trt代码生成引擎，但遇到显存分配问题。最终通过调整代码成功生成了引擎，但警告提示可能与cuBLAS版本不匹配。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

背景

最近解决了python版本为啥执行不了trtexec.exe生成的trt文件的bug，是环境里pytorch自带的cuda和trt的冲突，卸了重装了CPU版本就OK了。但是在我尝试解决的过程中出现了问题，环境有点玩坏了，之后用trtexec.exe生成的engine直接摆烂，输出的全是NaN。行吧，那我在python环境里生成吧。但是网上搜索的onnx2tensorrt代码基本上是7代之前的，我之前试了一些跑不出东西来。今天参考官方代码，记录以下tensorrt8在python环境下如何把onnx转为engine。

参考

github代码

简单流程

实际上用不到官方代码的那么多，只需要选一部分就行了：

import tensorrt as trt
import os

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


def get_engine(onnx_file_path, engine_file_path=""):
    """Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""

    def build_engine():
        """Takes an ONNX file and creates a TensorRT engine to run inference with"""
        with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
            EXPLICIT_BATCH
        ) as network, builder.create_builder_config() as config, trt.OnnxParser(
            network, TRT_LOGGER
        ) as parser, trt.Runtime(
            TRT_LOGGER
        ) as runtime:
            config.max_workspace_size = 1 << 32  # 4GB
            builder.max_batch_size = 1
            # Parse model file
            if not os.path.exists(onnx_file_path):
                print(
                    "ONNX file {} not found, please run yolov3_to_onnx.py first to generate it.".format(onnx_file_path)
                )
                exit(0)
            print("Loading ONNX file from path {}...".format(onnx_file_path))
            with open(onnx_file_path, "rb") as model:
                print("Beginning ONNX file parsing")
                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))
                    return None

            # # The actual yolov3.onnx is generated with batch size 64. Reshape input to batch size 1
            # network.get_input(0).shape = [1, 3, 608, 608]

            print("Completed parsing of ONNX file")
            print("Building an engine from file {}; this may take a while...".format(onnx_file_path))
            plan = builder.build_serialized_network(network, config)
            engine = runtime.deserialize_cuda_engine(plan)
            print("Completed creating Engine")
            with open(engine_file_path, "wb") as f:
                f.write(plan)
            return engine

    if os.path.exists(engine_file_path):
        # If a serialized engine exists, use it instead of building an engine.
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine()


def main():
    """Create a TensorRT engine for ONNX-based YOLOv3-608 and run inference."""

    # Try to load a previously generated YOLOv3-608 network graph in ONNX format:
    onnx_file_path = "model.onnx"
    engine_file_path = "model.trt"

    get_engine(onnx_file_path, engine_file_path)


if __name__ == "__main__":
    main()

摘抄下来之后我根据自己的模型删了一些东西，参考的这个模型的onnx好像batch是10，它自己就加了转为1的部分，我手头的模型文件已经是了，就删掉了。动态输入大小我没需要搞，就没看咋加。用pycharm试了下失败了，好像是显存分配有点问题，不知道设置的1<<32的大小是不是有点莽。但是我改到pycharm的终端敲代码执行成功了，就很抽象。

(mypytorch) PS F:\DeepStereo\AppleShow2> python onnx2trt.py      
onnx2trt.py:20: DeprecationWarning: Use set_memory_pool_limit instead.
  config.max_workspace_size = 1 << 32  # 4GB
Loading ONNX file from path G:\jupyter\Model_Zoo\resources_iter10_modify\crestereo_combined_iter10_240x320.onnx...
Beginning ONNX file parsing
[06/16/2022-16:59:16] [TRT] [W] onnx2trt_utils.cpp:365: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down t
o INT32.
Completed parsing of ONNX file
Building an engine from file G:\jupyter\Model_Zoo\resources_iter10_modify\crestereo_combined_iter10_240x320.onnx; this may take a while...
[06/16/2022-17:03:52] [TRT] [W] TensorRT was linked against cuBLAS/cuBLAS LT 11.8.0 but loaded cuBLAS/cuBLAS LT 11.3.1
[06/16/2022-17:07:42] [TRT] [W] TensorRT was linked against cuBLAS/cuBLAS LT 11.8.0 but loaded cuBLAS/cuBLAS LT 11.3.1
[06/16/2022-17:07:43] [TRT] [W] TensorRT was linked against cuBLAS/cuBLAS LT 11.8.0 but loaded cuBLAS/cuBLAS LT 11.3.1
Completed creating Engine