search

JingyunLiang / VRT

VRT: A Video Restoration Transformer (official repository)
https://arxiv.org/abs/2201.12288
Other
1.35k stars 126 forks source link

Exporting to onnx model problem #41

Open steven8274 opened 2 years ago

steven8274 commented 2 years ago

When I tried to export the 'VRT' model to onnx, the process was stopped by this error:

Traceback (most recent call last): File "main_test_vrt.py", line 424, in main() File "main_test_vrt.py", line 71, in main torch.onnx.export(model, # model being run File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/init.py", line 350, in export return utils.export( File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/utils.py", line 163, in export _export( File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/utils.py", line 1074, in _export graph, params_dict, torch_out = _model_to_graph( File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/utils.py", line 731, in _model_to_graph graph = _optimize_graph( File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/utils.py", line 308, in _optimize_graph graph = _C._jit_pass_onnx(graph, operator_export_type) File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/init.py", line 416, in _run_symbolic_function return utils._run_symbolic_function(*args, kwargs) File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/utils.py", line 1406, in _run_symbolic_function return symbolic_fn(g, *inputs, *attrs) File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/symbolic_helper.py", line 232, in wrapper return fn(g, args, kwargs) File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/symbolic_opset11.py", line 244, in pixel_shuffle return symbolic_helper._unimplemented("pixel_shuffle", "only support 4d input") File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/symbolic_helper.py", line 440, in _unimplemented _onnx_unsupported(f"{op}, {msg}") File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/symbolic_helper.py", line 444, in _onnx_unsupported raise RuntimeError( RuntimeError: Unsupported: ONNX export of operator pixel_shuffle, only support 4d input. Please feel free to request support or submit a pull request on PyTorch GitHub.

From the code I know we use pytorch 1.8.1 or higher version to support '5D PixelShuffle', however, ONNX only support 4d input.So, how can I fix this problem? Any help would be appreciated!