Revert "[ONNX] Remove deprecated export_to_pretty_string (pytorch#137790)"

This reverts commit d0628a7.

Reverted pytorch#137790 on behalf of https://github.com/facebook-github-bot due to Diff reverted internally ([comment](pytorch#137789 (comment)))

Author: pytorchmergebot

docs/source/onnx_torchscript.rst

@@ -697,6 +697,7 @@ Functions
 ^^^^^^^^^
 
 .. autofunction:: export
+.. autofunction:: export_to_pretty_string
 .. autofunction:: register_custom_op_symbolic
 .. autofunction:: unregister_custom_op_symbolic
 .. autofunction:: select_model_mode_for_export

test/onnx/test_pytorch_onnx_no_runtime.py

@@ -83,7 +83,7 @@ def forward(self, x):
 
         x = torch.ones(3, 3)
         f = io.BytesIO()
-        torch.onnx.export(AddmmModel(), x, f)
+        torch.onnx.export(AddmmModel(), x, f, verbose=False)
 
     def test_onnx_transpose_incomplete_tensor_type(self):
         # Smoke test to get us into the state where we are attempting to export
@@ -115,8 +115,7 @@ def foo(x):
 
         traced = torch.jit.trace(foo, (torch.rand([2])))
 
-        f = io.BytesIO()
-        torch.onnx.export(traced, (torch.rand([2]),), f)
+        torch.onnx.export_to_pretty_string(traced, (torch.rand([2]),))
 
     def test_onnx_export_script_module(self):
         class ModuleToExport(torch.jit.ScriptModule):
@@ -126,8 +125,7 @@ def forward(self, x):
                 return x + x
 
         mte = ModuleToExport()
-        f = io.BytesIO()
-        torch.onnx.export(mte, (torch.zeros(1, 2, 3),), f)
+        torch.onnx.export_to_pretty_string(mte, (torch.zeros(1, 2, 3),), verbose=False)
 
     @common_utils.suppress_warnings
     def test_onnx_export_func_with_warnings(self):
@@ -140,8 +138,9 @@ def forward(self, x):
                 return func_with_warning(x)
 
         # no exception
-        f = io.BytesIO()
-        torch.onnx.export(WarningTest(), torch.randn(42), f)
+        torch.onnx.export_to_pretty_string(
+            WarningTest(), torch.randn(42), verbose=False
+        )
 
     def test_onnx_export_script_python_fail(self):
         class PythonModule(torch.jit.ScriptModule):
@@ -162,7 +161,7 @@ def forward(self, x):
         mte = ModuleToExport()
         f = io.BytesIO()
         with self.assertRaisesRegex(RuntimeError, "Couldn't export Python"):
-            torch.onnx.export(mte, (torch.zeros(1, 2, 3),), f)
+            torch.onnx.export(mte, (torch.zeros(1, 2, 3),), f, verbose=False)
 
     def test_onnx_export_script_inline_trace(self):
         class ModuleToInline(torch.nn.Module):
@@ -180,8 +179,7 @@ def forward(self, x):
                 return y + y
 
         mte = ModuleToExport()
-        f = io.BytesIO()
-        torch.onnx.export(mte, (torch.zeros(1, 2, 3),), f)
+        torch.onnx.export_to_pretty_string(mte, (torch.zeros(1, 2, 3),), verbose=False)
 
     def test_onnx_export_script_inline_script(self):
         class ModuleToInline(torch.jit.ScriptModule):
@@ -200,8 +198,7 @@ def forward(self, x):
                 return y + y
 
         mte = ModuleToExport()
-        f = io.BytesIO()
-        torch.onnx.export(mte, (torch.zeros(1, 2, 3),), f)
+        torch.onnx.export_to_pretty_string(mte, (torch.zeros(1, 2, 3),), verbose=False)
 
     def test_onnx_export_script_module_loop(self):
         class ModuleToExport(torch.jit.ScriptModule):
@@ -215,8 +212,7 @@ def forward(self, x):
                 return x
 
         mte = ModuleToExport()
-        f = io.BytesIO()
-        torch.onnx.export(mte, (torch.zeros(1, 2, 3),), f)
+        torch.onnx.export_to_pretty_string(mte, (torch.zeros(1, 2, 3),), verbose=False)
 
     @common_utils.suppress_warnings
     def test_onnx_export_script_truediv(self):
@@ -228,8 +224,9 @@ def forward(self, x):
 
         mte = ModuleToExport()
 
-        f = io.BytesIO()
-        torch.onnx.export(mte, (torch.zeros(1, 2, 3, dtype=torch.float),), f)
+        torch.onnx.export_to_pretty_string(
+            mte, (torch.zeros(1, 2, 3, dtype=torch.float),), verbose=False
+        )
 
     def test_onnx_export_script_non_alpha_add_sub(self):
         class ModuleToExport(torch.jit.ScriptModule):
@@ -239,8 +236,7 @@ def forward(self, x):
                 return bs - 1
 
         mte = ModuleToExport()
-        f = io.BytesIO()
-        torch.onnx.export(mte, (torch.rand(3, 4),), f)
+        torch.onnx.export_to_pretty_string(mte, (torch.rand(3, 4),), verbose=False)
 
     def test_onnx_export_script_module_if(self):
         class ModuleToExport(torch.jit.ScriptModule):
@@ -251,8 +247,7 @@ def forward(self, x):
                 return x
 
         mte = ModuleToExport()
-        f = io.BytesIO()
-        torch.onnx.export(mte, (torch.zeros(1, 2, 3),), f)
+        torch.onnx.export_to_pretty_string(mte, (torch.zeros(1, 2, 3),), verbose=False)
 
     def test_onnx_export_script_inline_params(self):
         class ModuleToInline(torch.jit.ScriptModule):
@@ -282,8 +277,7 @@ def forward(self, x):
             torch.mm(torch.zeros(2, 3), torch.ones(3, 3)), torch.ones(3, 4)
         )
         self.assertEqual(result, reference)
-        f = io.BytesIO()
-        torch.onnx.export(mte, (torch.ones(2, 3),), f)
+        torch.onnx.export_to_pretty_string(mte, (torch.ones(2, 3),), verbose=False)
 
     def test_onnx_export_speculate(self):
         class Foo(torch.jit.ScriptModule):
@@ -318,10 +312,8 @@ def transpose(x):
         f1 = Foo(transpose)
         f2 = Foo(linear)
 
-        f = io.BytesIO()
-        torch.onnx.export(f1, (torch.ones(1, 10, dtype=torch.float),), f)
-        f = io.BytesIO()
-        torch.onnx.export(f2, (torch.ones(1, 10, dtype=torch.float),), f)
+        torch.onnx.export_to_pretty_string(f1, (torch.ones(1, 10, dtype=torch.float),))
+        torch.onnx.export_to_pretty_string(f2, (torch.ones(1, 10, dtype=torch.float),))
 
     def test_onnx_export_shape_reshape(self):
         class Foo(torch.nn.Module):
@@ -334,20 +326,17 @@ def forward(self, x):
                 return reshaped
 
         foo = torch.jit.trace(Foo(), torch.zeros(1, 2, 3))
-        f = io.BytesIO()
-        torch.onnx.export(foo, (torch.zeros(1, 2, 3)), f)
+        torch.onnx.export_to_pretty_string(foo, (torch.zeros(1, 2, 3)))
 
     def test_listconstruct_erasure(self):
         class FooMod(torch.nn.Module):
             def forward(self, x):
                 mask = x < 0.0
                 return x[mask]
 
-        f = io.BytesIO()
-        torch.onnx.export(
+        torch.onnx.export_to_pretty_string(
             FooMod(),
             (torch.rand(3, 4),),
-            f,
             add_node_names=False,
             do_constant_folding=False,
             operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK,
@@ -362,10 +351,13 @@ def forward(self, x):
                     retval += torch.sum(x[0:i], dim=0)
                 return retval
 
+        mod = DynamicSliceExportMod()
+
         input = torch.rand(3, 4, 5)
 
-        f = io.BytesIO()
-        torch.onnx.export(DynamicSliceExportMod(), (input,), f, opset_version=10)
+        torch.onnx.export_to_pretty_string(
+            DynamicSliceExportMod(), (input,), opset_version=10
+        )
 
     def test_export_dict(self):
         class DictModule(torch.nn.Module):
@@ -376,12 +368,10 @@ def forward(self, x_in: torch.Tensor) -> Dict[str, torch.Tensor]:
         mod = DictModule()
         mod.train(False)
 
-        f = io.BytesIO()
-        torch.onnx.export(mod, (x_in,), f)
+        torch.onnx.export_to_pretty_string(mod, (x_in,))
 
         with self.assertRaisesRegex(RuntimeError, r"DictConstruct.+is not supported."):
-            f = io.BytesIO()
-            torch.onnx.export(torch.jit.script(mod), (x_in,), f)
+            torch.onnx.export_to_pretty_string(torch.jit.script(mod), (x_in,))
 
     def test_source_range_propagation(self):
         class ExpandingModule(torch.nn.Module):
@@ -507,11 +497,11 @@ def forward(self, box_regression: Tensor, proposals: List[Tensor]):
         proposal = [torch.randn(2, 4), torch.randn(2, 4)]
 
         with self.assertRaises(RuntimeError) as cm:
-            f = io.BytesIO()
+            onnx_model = io.BytesIO()
             torch.onnx.export(
                 model,
                 (box_regression, proposal),
-                f,
+                onnx_model,
             )
 
     def test_initializer_sequence(self):
@@ -647,7 +637,7 @@ def forward(self, x):
 
         x = torch.randn(1, 2, 3, requires_grad=True)
         f = io.BytesIO()
-        torch.onnx.export(Model(), (x,), f)
+        torch.onnx.export(Model(), x, f)
         model = onnx.load(f)
         model.ir_version = 0
 
@@ -754,7 +744,7 @@ def forward(self, x):
 
         f = io.BytesIO()
         with warnings.catch_warnings(record=True):
-            torch.onnx.export(MyDrop(), (eg,), f)
+            torch.onnx.export(MyDrop(), (eg,), f, verbose=False)
 
     def test_pack_padded_pad_packed_trace(self):
         from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
@@ -801,7 +791,7 @@ def forward(self, x, seq_lens):
         self.assertEqual(grad, grad_traced)
 
         f = io.BytesIO()
-        torch.onnx.export(m, (x, seq_lens), f)
+        torch.onnx.export(m, (x, seq_lens), f, verbose=False)
 
     # Suppression: ONNX warns when exporting RNNs because of potential batch size mismatch.
     @common_utils.suppress_warnings
@@ -861,7 +851,7 @@ def forward(self, x, seq_lens):
             self.assertEqual(grad, grad_traced)
 
             f = io.BytesIO()
-            torch.onnx.export(m, (x, seq_lens), f)
+            torch.onnx.export(m, (x, seq_lens), f, verbose=False)
 
     def test_pushpackingpastrnn_in_peephole_create_own_gather_input(self):
         from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
@@ -941,8 +931,7 @@ class Mod(torch.nn.Module):
             def forward(self, x, w):
                 return torch.matmul(x, w).detach()
 
-        f = io.BytesIO()
-        torch.onnx.export(Mod(), (torch.rand(3, 4), torch.rand(4, 5)), f)
+        torch.onnx.export_to_pretty_string(Mod(), (torch.rand(3, 4), torch.rand(4, 5)))
 
     def test_aten_fallback_must_fallback(self):
         class ModelWithAtenNotONNXOp(torch.nn.Module):
@@ -1099,12 +1088,12 @@ def sym_scatter_max(g, src, index, dim, out, dim_size):
         torch.onnx.register_custom_op_symbolic(
             "torch_scatter::scatter_max", sym_scatter_max, 1
         )
-        f = io.BytesIO()
         with torch.no_grad():
             torch.onnx.export(
                 m,
                 (src, idx),
-                f,
+                "mymodel.onnx",
+                verbose=False,
                 opset_version=13,
                 custom_opsets={"torch_scatter": 1},
                 do_constant_folding=True,
@@ -1187,7 +1176,7 @@ def forward(self, x):
         model = Net(C).cuda().half()
         x = torch.randn(N, C).cuda().half()
         f = io.BytesIO()
-        torch.onnx.export(model, (x,), f, opset_version=14)
+        torch.onnx.export(model, x, f, opset_version=14)
         onnx_model = onnx.load_from_string(f.getvalue())
         const_node = [n for n in onnx_model.graph.node if n.op_type == "Constant"]
         self.assertNotEqual(len(const_node), 0)

torch/onnx/__init__.py

@@ -30,6 +30,7 @@
     "JitScalarType",
     # Public functions
     "export",
+    "export_to_pretty_string",
     "is_in_onnx_export",
     "select_model_mode_for_export",
     "register_custom_op_symbolic",
@@ -67,6 +68,7 @@
 from .utils import (
     _run_symbolic_function,
     _run_symbolic_method,
+    export_to_pretty_string,
     is_in_onnx_export,
     register_custom_op_symbolic,
     select_model_mode_for_export,

torch/onnx/utils.py

@@ -35,6 +35,7 @@
     "model_signature",
     "warn_on_static_input_change",
     "unpack_quantized_tensor",
+    "export_to_pretty_string",
     "unconvertible_ops",
     "register_custom_op_symbolic",
     "unregister_custom_op_symbolic",
@@ -1139,6 +1140,84 @@ def _model_to_graph(
     return graph, params_dict, torch_out
 
 
+@torch._disable_dynamo
+@_deprecation.deprecated("2.5", "the future", "use onnx.printer.to_text() instead")
+def export_to_pretty_string(
+    model,
+    args,
+    export_params=True,
+    verbose=False,
+    training=_C_onnx.TrainingMode.EVAL,
+    input_names=None,
+    output_names=None,
+    operator_export_type=_C_onnx.OperatorExportTypes.ONNX,
+    export_type=None,
+    google_printer=False,
+    opset_version=None,
+    keep_initializers_as_inputs=None,
+    custom_opsets=None,
+    add_node_names=True,
+    do_constant_folding=True,
+    dynamic_axes=None,
+):
+    """Similar to :func:`export`, but returns a text representation of the ONNX model.
+
+    Only differences in args listed below. All other args are the same
+    as :func:`export`.
+
+    Args:
+        add_node_names (bool, default True): Whether or not to set
+            NodeProto.name. This makes no difference unless
+            ``google_printer=True``.
+        google_printer (bool, default False): If False, will return a custom,
+            compact representation of the model. If True will return the
+            protobuf's `Message::DebugString()`, which is more verbose.
+
+    Returns:
+        A UTF-8 str containing a human-readable representation of the ONNX model.
+    """
+    if opset_version is None:
+        opset_version = _constants.ONNX_DEFAULT_OPSET
+    if custom_opsets is None:
+        custom_opsets = {}
+    GLOBALS.export_onnx_opset_version = opset_version
+    GLOBALS.operator_export_type = operator_export_type
+
+    with exporter_context(model, training, verbose):
+        val_keep_init_as_ip = _decide_keep_init_as_input(
+            keep_initializers_as_inputs, operator_export_type, opset_version
+        )
+        val_add_node_names = _decide_add_node_names(
+            add_node_names, operator_export_type
+        )
+        val_do_constant_folding = _decide_constant_folding(
+            do_constant_folding, operator_export_type, training
+        )
+        args = _decide_input_format(model, args)
+        graph, params_dict, torch_out = _model_to_graph(
+            model,
+            args,
+            verbose,
+            input_names,
+            output_names,
+            operator_export_type,
+            val_do_constant_folding,
+            training=training,
+            dynamic_axes=dynamic_axes,
+        )
+
+        return graph._pretty_print_onnx(  # type: ignore[attr-defined]
+            params_dict,
+            opset_version,
+            False,
+            operator_export_type,
+            google_printer,
+            val_keep_init_as_ip,
+            custom_opsets,
+            val_add_node_names,
+        )
+
+
 @_deprecation.deprecated("2.5", "the future", "avoid using this function")
 def unconvertible_ops(
     model,