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

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 
    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)
 

你可能感兴趣的:(计算机视觉,深度学习)