BUGFix:onnx -> TensorRT转换过程失败

先附上相关的onnx2trt的部分代码:

def onnx2trt(onnx_path):
    logger = trt.Logger(trt.Logger.ERROR)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)
    parser.parse_from_file(onnx_path)
    config = builder.create_builder_config()
    config.max_workspace_size=max_workspace_size
    config.set_flag(trt.BuilderFlag.FP16)
    op = builder.create_optimization_profile()
    # op.set_shape('model0/input', (1, )+shape, (batch_size[0], )+shape, (batch_size[1], )+shape)
    op.set_shape(network.get_input(0).name, (min_batch_size, )+input_shape, (opt_batch_size, )+input_shape, (max_batch_size, )+input_shape)
    config.add_optimization_profile(op)
    engine = builder.build_engine(network, config)
    # trt_path = onnx_path.replace('/onnx/', '/trt/').replace('.onnx', '.plan')
    trt_path = onnx_path.replace('.onnx', '.plan')
    with open(trt_path,'wb') as f:
        f.write(engine.serialize())

在onnx转换TensorRT的过程中,提示15行代码有错误:

config.max_workspace_size=max_workspace_size

其中,max_workspace_size = 1<<30

首先单位是字节,比如 builder.max_workspace_size = 1<< 30 就是 2^30 bytes 即 1 GB。

它的作用是给出模型中任一层能使用的内存上限。运行时,每一层需要多少内存系统分配多少,并不是每次都分 1 GB,但不会超过 1 GB。

具体报错信息如下:

TypeError: deserialize_cuda_engine(): incompatible function arguments. The following argument types are supported:
   1. (self: tensorrt.tensorrt.Runtime, serialized_engine: buffer) -> tensorrt.tensorrt.ICudaEngine

Invoked with: <tensorrt.tensorrt.Runtime object at 0x7feecb3c6530>, None

上面这错误可能是由于max_workspace_size分配不够导致的错误,可试着将30放大,但是我这里不管用;

原因是构建nvidia-docker时候,设置 --shm-size =32,共享内存的太小,不支持onnx-TensorRT的操作,这里修改为64,问题得以解决;

有问题随时交流,欢迎一键三连~

参考:

https://www.cnblogs.com/mrlonely2018/p/14841562.html

相关推荐
手插口袋谁也不爱♡8 分钟前
本地化语音识别CapsWriter结合内网穿透远程会议录音秒变文字稿
人工智能·语音识别
deephub18 分钟前
FANformer:融合傅里叶分析网络的大语言模型基础架构
人工智能·语言模型·傅立叶分析
飞哥数智坊25 分钟前
Cursor实战:1小时集成天地图
人工智能·cursor
ianozo30 分钟前
[GHCTF 2025]UPUPUP【.htaccess绕过 XBM/WBMP】
图像处理·人工智能
大囚长1 小时前
deepseek+ansible实现AI自动化集群部署
人工智能·自动化·ansible
程序边界1 小时前
AI+游戏开发:如何用 DeepSeek 打造高性能贪吃蛇游戏
人工智能·游戏
CodeJourney.1 小时前
光储直流微电网:能源转型的关键力量
数据库·人工智能·算法·能源
艾思科蓝 AiScholar1 小时前
【 IEEE出版 | 快速稳定EI检索 | 往届已EI检索】2025年储能及能源转换国际学术会议(ESEC 2025)
人工智能·计算机网络·自然语言处理·数据挖掘·自动化·云计算·能源
Fulima_cloud1 小时前
智慧锂电:开启能源新时代的钥匙
大数据·人工智能·物联网
GUOYUGRA1 小时前
高纯氢能源在线监测分析系统组成和作用
人工智能·算法·机器学习