OpenDriveLab / UniAD

[CVPR 2023 Best Paper Award] Planning-oriented Autonomous Driving
Apache License 2.0
3.53k stars 400 forks source link

[Export][ONNX] LiDARInstance3DBoxes #198

Open Lewis-Lu opened 2 months ago

Lewis-Lu commented 2 months ago

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 commented 2 months ago

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

Lewis-Lu commented 2 months ago

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

Jian-danai commented 3 weeks ago

See https://github.com/OpenDriveLab/UniAD/issues/81