[Export][ONNX] LiDARInstance3DBoxes

[Lewis-Lu]

Hi there,

I'm trying to export motion_head to ONNX, as the track_bbox_results contains LiDARInstance3DBoxes, thus the torch.onnx.tracer failed with the following log:

Restoring the weights from ckpts/uniad_base_e2e.pth.

export motion_head, #params = 13.77M

Traceback (most recent call last):
File "tools/export.py", line 121, in
main()
File "tools/export.py", line 113, in main
torch.onnx.export(model,
File "/data_ssd/hlu/miniconda3/envs/bevformer/lib/python3.8/site-packages/torch/onnx/init.py", line 275, in export
return utils.export(model, args, f, export_params, verbose, training,
File "/data_ssd/hlu/miniconda3/envs/bevformer/lib/python3.8/site-packages/torch/onnx/utils.py", line 88, in export
_export(model, args, f, export_params, verbose, training, input_names, output_names,
File "/data_ssd/hlu/miniconda3/envs/bevformer/lib/python3.8/site-packages/torch/onnx/utils.py", line 689, in _export
_model_to_graph(model, args, verbose, input_names,
File "/data_ssd/hlu/miniconda3/envs/bevformer/lib/python3.8/site-packages/torch/onnx/utils.py", line 458, in _model_to_graph
graph, params, torch_out, module = _create_jit_graph(model, args,
File "/data_ssd/hlu/miniconda3/envs/bevformer/lib/python3.8/site-packages/torch/onnx/utils.py", line 422, in _create_jit_graph
graph, torch_out = _trace_and_get_graph_from_model(model, args)
File "/data_ssd/hlu/miniconda3/envs/bevformer/lib/python3.8/site-packages/torch/onnx/utils.py", line 373, in _trace_and_get_graph_from_model
torch.jit._get_trace_graph(model, args, strict=False, _force_outplace=False, _return_inputs_states=True)
File "/data_ssd/hlu/miniconda3/envs/bevformer/lib/python3.8/site-packages/torch/jit/_trace.py", line 1160, in _get_trace_graph
outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
File "/data_ssd/hlu/miniconda3/envs/bevformer/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
return forward_call(*input, **kwargs)
File "/data_ssd/hlu/miniconda3/envs/bevformer/lib/python3.8/site-packages/torch/jit/_trace.py", line 95, in forward
in_vars, in_desc = _flatten(args)
RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: LiDARInstance3DBoxes

1. I can manually decode the data store in 3d-bbox, then tracer could work, since the forward() args would be torch.Tensor then. However, it takes much effort for exporting the part to ONNX.
2. Can we plugin some code in torch.jit module for support this? or any MMDeploy code snippets might help ? I'm not familiar with mmdeploy stuff. So might somebody could give this issue a hint.

Best,
Lewis

[Lewis-Lu]

@YTEP-ZHI could you please comment on this ? Many Thanks in advance

[Lewis-Lu]

Some more context, I'm looking into open-mmlab/mmdeploy#1662 for some hints as it seems has supported InstanceData for exporting to ONNX.

[Jian-danai]

See #81