Introduce torch.onnx.dynamo_export API
This is the first phase of the new ONNX exporter API for exporting from
TorchDynamo and FX, and represents the beginning of a new era for exporting
ONNX from PyTorch.
The API here is a starting point upon which we will layer more capability
and expressiveness in subsequent phases. This first phase introduces the
following into `torch.onnx`:
dynamo_export(
model: torch.nn.Module,
/,
*model_args,
export_options: Optional[ExportOptions] = None,
**model_kwargs,
) -> ExportOutput
class ExportOptions:
opset_version: Optional[int] = None
dynamic_shapes: Optional[bool] = None
logger: Optional[logging.Logger] = None
class ExportOutputSerializer(Protocol):
def serialize(
self,
export_output: ExportOutput,
destination: io.BufferedIOBase
) -> None:
...
class ExportOutput:
export_options: ExportOptions
model_proto: onnx.ModelProto
def save(
self,
destination: Union[str, io.BufferedIOBase],
*,
serializer: Optional[ExportOutputSerializer] = None
) -> None:
...
Co-authored-by: Bowen Bao <bowbao@microsoft.com>
Co-authored-by: Aaron Bockover <abock@microsoft.com>
