TensorRT 5/6/7 加载onnx模型

最新推荐文章于 2025-06-16 11:23:40 发布

cx媛

最新推荐文章于 2025-06-16 11:23:40 发布

阅读量5.9k

点赞数 5

CC 4.0 BY-SA版权

分类专栏：深度学习 tensorrt 文章标签：深度学习

本文链接：https://blog.youkuaiyun.com/qq_23116521/article/details/107617442

本文介绍了如何在TensorRT 5/6及7中加载和转换ONNX模型。重点讲述了在TensorRT 7中由于对explicit_batch的要求，需要设置`explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)`，否则会出现错误。此外，还强调了在TRT7中必须指定输入尺寸，否则网络验证会失败。作者分享了遇到的问题及解决方案，包括int8量化时动态batch_size导致的错误，并给出了处理方法。

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

TensorRT5/6加载并转换onnx:

def onnx_2_trt(onnx_model_name,trt_model_name):
    with trt.Builder(G_LOGGER) as builder, builder.create_network() as network, 
    trt.OnnxParser(network, G_LOGGER) as parser:
        builder.max_batch_size = 1024
        builder.max_workspace_size = 2 << 30

        print('Loading ONNX file from path {}...'.format(onnx_model_name))
        with open(onnx_model_name, 'rb') as model:
            print('Beginning ONNX file parsing')
            b = parser.parse(model.read())

        if 1:
            print('Completed parsing of ONNX file')
            print('Building an engine from file {}; this may take a while...'.format(args.onnx_model))

            ####
            #builder.int8_mode = True
            #builder.int8_calibrator = calib
            builder.fp16_mode = True
            ####
            print("layers:",network.num_layers)
            network.mark_output(network.get_layer(network.num_layers - 1).get_output(0))//有的模型需要，有的模型在转onnx的之后已经指定了，就不需要这行

            engine = builder.build_cuda_engine(network)
            print(engine)
            print("Completed creating Engine")
            with open(trt_model_name, "wb") as f:
                f.write(engine.serialize())
            return engine

        else:
            print('Number of errors: {}'.format(parser.num_errors))
            error = parser.get_error(0) #