pytorch
51ce0223 - [ONNX] Support converting fx graph with symbolic shape to ONNX (#96350)

Commit
1 year ago
[ONNX] Support converting fx graph with symbolic shape to ONNX (#96350) ~~Need https://github.com/microsoft/onnx-script/pull/484~~ Support dynamic export on fx-ONNX exporter. Essentially, we set inputs size and nodes all dynamic in torchscript, and leverage on `aten::sym_size` to catch dynamic size between each Op. 1. Add `dynamic_axes` switch between symbolic tracing (dynamic sizes) and fake mode (static). Set it to default True, as most of our tests are happy with sumbolic tracing. Except GPT2 stays with fake mode with error: https://github.com/microsoft/onnx-script/issues/523 2. Add test_fx_dynamic_onnruntime.py to test on some addhoc we have from old exporter. This can be removed once tests are integrated with https://github.com/pytorch/pytorch/pull/96479 3. Since `aten::sym_size` are operated with built-in function, a built-in function mapping is added to support SymFloat/SymInt. (FIXME: https://github.com/pytorch/pytorch/issues/97201). sym_size output value is also fx.Node, and can be found in `fx_name_to_onnxscipt_value`, so it's operation stays the same as other ONNX ops in ONNX graph. 4. Fully deprecated FakeTensorProp as make_fx() should provide all node meta info. 5. Put complicated fx.Node related ArgumentType into _type_utils.py Pull Request resolved: https://github.com/pytorch/pytorch/pull/96350 Approved by: https://github.com/wschin, https://github.com/justinchuby
Author
Committer
Parents
Loading