Pytorch模型转换,pth->onnx->trt(TensorRT engine)

本文档详细解析了在将PyTorch模型(pth格式)转换为ONNX格式过程中遇到的常见问题,包括模型输入输出不支持字典类型、tensor.size(dim)操作导致的Gather[0]问题及求余算符在导出ONNX时不支持等,并提供了相应的解决策略。
部署运行你感兴趣的模型镜像
pth->onnx常见问题


##模型输入输出不支持字典
在使用torch.onnx导出onnx格式时,模型的输入和输出都不支持字典类型的数据结构。


**解决方法:**
此时,可以将字典数据结构换为torch.onnx支持的列表或者元组。例如:
heads {'hm': 1, 'wh': 2, 'hps': 34, 'reg': 2, 'hm_hp': 17, 'hp_offset': 2}
Centerpose中的字典在导出onnx时将会报错,可以将字典拆为两个列表,然后模型代码中做相应修改。


##tensor.size(dim)操作导致的转出的onnx操作中带有Gather[0]
如果在网络中存在类似tensor.size(dim)或者tensor.shape[1]这种索引操作,就会在转出的onnx模型中产生出Gather[0]操作,而这个操作,在TensorRT中是不支持的。


**解决方法:**
尽量避免类似的操作,代码中不要出现与网络操作无关的tensor,一些尺寸之类的常数全部写成普通变量。


##求余算符在导出onnx时不支持
曾经在导出求余算符‘%’时出现如下报错,是因为此运算符为ATen操作符,但是未在torch/onnx/symbolic_opset9.py中定义。
/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py:562: UserWarning: ONNX export failed on ATen operator remainder because torch.onnx.symbolic_opset9.remainder does not exist
  .format(op_name, opset_version, op_name))
Traceback (most recent call last):
  File "transfer_centernet_dlav0_to_onnx2.py", line 65, in <module>
    onnx_module = torch.onnx.export(model,img,'multipose_dlav0_1x_modellast_process.onnx',verbose=True)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/__init__.py", line 132, in export
    strip_doc_string, dynamic_axes)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py", line 64, in export
    example_outputs=example_outputs, strip_doc_string=strip_doc_string, dynamic_axes=dynamic_axes)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py", line 329, in _export
    _retain_param_name, do_constant_folding)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py", line 225, in _model_to_graph
    _disable_torch_constant_prop=_disable_torch_constant_prop)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py", line 127, in _optimize_graph
    graph = torch._C._jit_pass_onnx(graph, operator_export_type)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/__init__.py", line 163, in _run_symbolic_function
    return utils._run_symbolic_function(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py", line 563, in _run_symbolic_function
    op_fn = sym_registry.get_registered_op(op_name, '', opset_version)
  File "/usr/local/lib/python3.6/dist-packages/torch/onnx/symbolic_registry.py", line 91, in get_registered_op
    return _registry[(domain, version)][opname]
KeyError: 'remainder'


**解决方法:**
在torch/onnx/symbolic_opset9.py中补上自己的remainder的定义,添加的代码如下
@parse_args( 'v', 'v')
def remainder(g,input,division):
    return g.op("Remainder",input,division)

 

您可能感兴趣的与本文相关的镜像

PyTorch 2.5

PyTorch 2.5

PyTorch
Cuda

PyTorch 是一个开源的 Python 机器学习库,基于 Torch 库,底层由 C++ 实现,应用于人工智能领域,如计算机视觉和自然语言处理

<think>我们正在处理用户关于模型转换的问题。用户询问如何将ONNXPyTorch模型转换TensorRT引擎。根据提供的引用,我们可以总结出两种主要方法: 1. 直接使用TensorRT的trtexec工具将ONNX模型转换TensorRT引擎(引用[1][2][4]提到) 2. 通过Python API进行转换(引用[2][4]提到) 同时,用户提到他们有一个训练好的模型ONNXPyTorch格式)。注意,对于PyTorch模型,通常需要先转换ONNX,然后再转换TensorRT(引用[3][4])。 因此,我们的回答应该分为两部分: 第一部分:如何将PyTorch模型转换ONNX模型(如果用户需要的话) 第二部分:如何将ONNX模型转换TensorRT引擎(通过两种方式:trtexec命令行工具和Python API) 另外,引用中提到了转换过程中的一些注意事项,比如验证ONNX模型结构、处理转换中的常见错误等,我们也应该提及。 由于用户特别询问了ONNXPyTorch模型转换TensorRT引擎,我们可以这样组织: 步骤1:PyTorch模型ONNX(如果已有ONNX模型可跳过) 步骤2:ONNX模型TensorRT引擎(两种方法) 同时,我们还需要注意: -转换过程中,设置动态batch size(引用[3]提到) - 使用合适的opset版本(引用[4]中使用了opset_version=11) 最后,按照用户要求,在回答末尾生成相关问题。 注意:引用标识要自然地添加在相关内容的末尾。 开始构建回答:</think>### 如何将 ONNXPyTorch 模型转换TensorRT 引擎 以下是基于引用内容的完整转换流程,涵盖 PyTorchONNXTensorRT 的核心步骤和关键注意事项: --- #### **一、PyTorch 模型ONNX 模型** [^1][^4] ```python import torch # 1. 加载预训练模型并切换为推理模式 model = YourModel() # 替换为实际模型 model.load_state_dict(torch.load("model.pth")) model.eval() # ⚠️ 必须关闭 dropout 和 batch norm 的随机性 # 2. 准备虚拟输入(需匹配实际输入尺寸) dummy_input = torch.randn(1, 3, 224, 224) # 示例:(batch, channel, height, width) # 3. 导出 ONNX(关键参数说明) torch.onnx.export( model, dummy_input, "model.onnx", # 输出文件名 export_params=True, # 包含模型权重 opset_version=13, # 推荐 ≥11(兼容 TensorRT) input_names=['input'], # 输入节点名 output_names=['output'], # 输出节点名 dynamic_axes={ 'input': {0: 'batch'}, # ⚠️ 启用动态 batch 维度[^3] 'output': {0: 'batch'} } ) ``` **关键验证**: - 使用 `onnx.checker.check_model("model.onnx")` 验证模型结构 - 通过 Netron 可视化工具检查算子兼容性[^2] --- #### **二、ONNX 模型TensorRT 引擎** ##### **方法 1:使用 `trtexec` 命令行工具** [^1][^2] ```bash # 基础转换(FP32 精度) trtexec --onnx=model.onnx --saveEngine=model.engine # 常用优化参数: --fp16 # 启用 FP16 加速(推荐 Volta+ 显卡) --int8 # 启用 INT8 量化(需校准数据集) --workspace=4096 # 分配显存空间(单位 MB) --minShapes=input:1x3x224x224 # 动态尺寸下限 --optShapes=input:8x3x224x224 # 最优性能尺寸 --maxShapes=input:32x3x224x224 # 动态尺寸上限 ``` ##### **方法 2:Python API 转换** [^2][^4] ```python import tensorrt as trt # 1. 创建构建器和配置 logger = trt.Logger(trt.Logger.INFO) builder = trt.Builder(logger) config = builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2 << 30) # 2GB 显存 # 2. 解析 ONNX 模型 network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, logger) with open("model.onnx", "rb") as f: parser.parse(f.read()) # 3. 设置优化配置(可选) config.set_flag(trt.BuilderFlag.FP16) # 启用 FP16 # config.set_flag(trt.BuilderFlag.INT8) # 启用 INT8(需额外校准) # 4. 构建引擎并序列化 engine = builder.build_serialized_network(network, config) with open("model.engine", "wb") as f: f.write(engine) ``` --- #### **三、常见问题与解决方案** [^2] | 错误类型 | 原因 | 解决方法 | |------------------------|-----------------------------|----------------------------| | `Unsupported ONNX op` | 算子未在 TensorRT 中实现 | 替换为兼容算子或自定义插件 | | `Shape tensor error` | 动态尺寸配置错误 | 检查 `dynamic_axes` 设置一致性 | | `FP16 overflow` | 数值范围溢出 | 添加 `LayerNorm` 或降低精度 | | `INT8 calibration fail`| 校准数据分布与真实数据不符 | 使用代表性校准数据集 | > **最佳实践**: > - 动态 Batch:ONNX 导出时声明 `dynamic_axes`,TensorRT 构建时指定 `min/opt/maxShapes`[^3] > - 精度选择:优先尝试 `FP16`,关键任务系统用 `FP32`,边缘设备用 `INT8` > - 性能验证:使用 `trtexec --loadEngine=model.engine --shapes=input:8x3x224x224` 测试吞吐量 --- #### **四、推理部署示例(Python)** ```python import tensorrt as trt import pycuda.driver as cuda # 加载引擎 with open("model.engine", "rb") as f, trt.Runtime(trt.Logger(trt.Logger.INFO)) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) # 创建执行上下文 context = engine.create_execution_context() # 分配 GPU 内存 inputs, outputs, bindings = [], [], [] stream = cuda.Stream() for binding in engine: size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size dtype = trt.nptype(engine.get_binding_dtype(binding)) # 分配内存并绑定 mem = cuda.mem_alloc(size * dtype.itemsize) bindings.append(int(mem)) if engine.binding_is_input(binding): inputs.append({'mem': mem, 'size': size, 'dtype': dtype}) else: outputs.append({'mem': mem, 'size': size, 'dtype': dtype}) # 执行推理 input_data = np.random.random((8, 3, 224, 224)).astype(np.float32) # 示例输入 cuda.memcpy_htod_async(inputs[0]['mem'], input_data, stream) context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) output_data = np.empty(outputs[0]['size'], dtype=outputs[0]['dtype']) cuda.memcpy_dtoh_async(output_data, outputs[0]['mem'], stream) stream.synchronize() ``` ---
评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值