codekingpro/portable-devtools

1# -------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation.  All rights reserved.
3# Licensed under the MIT License.
4# --------------------------------------------------------------------------
5
6import torch
7from torch._C._onnx import OperatorExportTypes
8
9TrainingMode = torch.onnx.TrainingMode
10from packaging.version import Version  # noqa: E402
11
12
13def torch_onnx_export(
14    model,
15    args,
16    f,
17    export_params=True,
18    verbose=False,
19    training=TrainingMode.EVAL,
20    input_names=None,
21    output_names=None,
22    operator_export_type=OperatorExportTypes.ONNX,
23    opset_version=None,
24    _retain_param_name=None,
25    do_constant_folding=True,
26    example_outputs=None,
27    strip_doc_string=None,
28    dynamic_axes=None,
29    keep_initializers_as_inputs=None,
30    custom_opsets=None,
31    enable_onnx_checker=None,
32    use_external_data_format=None,
33    export_modules_as_functions=False,
34):
35    if Version(torch.__version__) >= Version("1.11.0"):
36        torch.onnx.export(
37            model=model,
38            args=args,
39            f=f,
40            export_params=export_params,
41            verbose=verbose,
42            training=training,
43            input_names=input_names,
44            output_names=output_names,
45            operator_export_type=operator_export_type,
46            opset_version=opset_version,
47            do_constant_folding=do_constant_folding,
48            dynamic_axes=dynamic_axes,
49            keep_initializers_as_inputs=keep_initializers_as_inputs,
50            custom_opsets=custom_opsets,
51            export_modules_as_functions=export_modules_as_functions,
52            dynamo=False,
53        )
54    else:
55        torch.onnx.export(
56            model=model,
57            args=args,
58            f=f,
59            export_params=export_params,
60            verbose=verbose,
61            training=training,
62            input_names=input_names,
63            output_names=output_names,
64            operator_export_type=operator_export_type,
65            opset_version=opset_version,
66            _retain_param_name=_retain_param_name,
67            do_constant_folding=do_constant_folding,
68            example_outputs=example_outputs,
69            strip_doc_string=strip_doc_string,
70            dynamic_axes=dynamic_axes,
71            keep_initializers_as_inputs=keep_initializers_as_inputs,
72            custom_opsets=custom_opsets,
73            enable_onnx_checker=enable_onnx_checker,
74            use_external_data_format=use_external_data_format,
75        )
76