Setting automatic default selection for ONNX IR v4 semantics in ONNX export API (pytorch#26146)

Summary:
This is a follow-up PR for pytorch#23284. In that PR we had removed changing the default behavior for `keep_initializers_as_input` argument to the export API. With this PR we are enabling that change in that if `keep_initializers_as_input` is not specified then value/behavior for this argument is chosen automatically depending on whether the export type is ONNX or not.

This was part of the earlier PR was removed for further review. The test points have also been updated.

This change may fail some internal tests which may require explicitly setting `keep_initializers_as_input=True` to preserve old behavior.
Pull Request resolved: pytorch#26146

Reviewed By: hl475

Differential Revision: D17369677

Pulled By: houseroad

fbshipit-source-id: 2aec2cff50d215714ee8769505ef24d2b7865a11

Author: Spandan Tiwari

test/onnx/test_operators.py

@@ -247,7 +247,7 @@ def test_batchnorm(self):
 
     def test_batchnorm_onnx_irv4(self):
         x = torch.ones(2, 2, 2, 2, requires_grad=True)
-        self.assertONNX(nn.BatchNorm2d(2), x, keep_initializers_as_inputs=False)
+        self.assertONNX(nn.BatchNorm2d(2), x)
 
     def test_batchnorm_1d(self):
         x = torch.ones(2, 2, requires_grad=True)
@@ -263,7 +263,7 @@ def test_conv(self):
 
     def test_conv_onnx_irv4(self):
         x = torch.ones(20, 16, 50, 40, requires_grad=True)
-        self.assertONNX(nn.Conv2d(16, 13, 3, bias=False), x, keep_initializers_as_inputs=False)
+        self.assertONNX(nn.Conv2d(16, 13, 3, bias=False), x)
 
     def test_conv_variable_length(self):
         x = torch.ones(5, 3, 6, 6, requires_grad=True)

torch/onnx/__init__.py

@@ -31,7 +31,7 @@ def export(model, args, f, export_params=True, verbose=False, training=False,
            input_names=None, output_names=None, aten=False, export_raw_ir=False,
            operator_export_type=None, opset_version=None, _retain_param_name=True,
            do_constant_folding=False, example_outputs=None, strip_doc_string=True, 
-           dynamic_axes=None, keep_initializers_as_inputs=True):
+           dynamic_axes=None, keep_initializers_as_inputs=None):
     r"""
     Export a model into ONNX format.  This exporter runs your model
     once in order to get a trace of its execution to be exported;
@@ -123,12 +123,16 @@ def export(model, args, f, export_params=True, verbose=False, training=False,
 
                 (c). MIXED MODE OF (a) and (b)
                     dynamic_axes = {'input_1':[0, 2, 3], 'input_2':{0:'batch'}, 'output':[0,1]}
-        keep_initializers_as_inputs (bool, default True): If True, all the initializers
+        keep_initializers_as_inputs (bool, default None): If True, all the initializers
             (typically corresponding to parameters) in the exported graph will also be 
             added as inputs to the graph. If False, then initializers are not added as
             inputs to the graph, and only the non-parameter inputs are added as inputs.
             This may allow for better optimizations (such as constant folding etc.) by
-            backends/runtimes that execute these graphs.
+            backends/runtimes that execute these graphs. If unspecified (default None),
+            then the behavior is chosen automatically as follows. If operator_export_type
+            is OperatorExportTypes.ONNX, the behavior is equivalent to setting this
+            argument to False. For other values of operator_export_type, the behavior is
+            equivalent to setting this argument to True.
     """
 
     from torch.onnx import utils

torch/onnx/utils.py

@@ -49,7 +49,7 @@ def export(model, args, f, export_params=True, verbose=False, training=False,
            input_names=None, output_names=None, aten=False, export_raw_ir=False,
            operator_export_type=None, opset_version=None, _retain_param_name=True,
            do_constant_folding=False, example_outputs=None, strip_doc_string=True,
-           dynamic_axes=None, keep_initializers_as_inputs=True):
+           dynamic_axes=None, keep_initializers_as_inputs=None):
     if aten or export_raw_ir:
         assert operator_export_type is None
         assert aten ^ export_raw_ir
@@ -276,7 +276,7 @@ def export_to_pretty_string(model, args, f, export_params=True, verbose=False, t
                             operator_export_type=None, export_type=ExportTypes.PROTOBUF_FILE,
                             example_outputs=None, propagate=False, google_printer=False,
                             opset_version=None, _retain_param_name=True,
-                            keep_initializers_as_inputs=True):
+                            keep_initializers_as_inputs=None):
     if aten or export_raw_ir:
         assert operator_export_type is None
         assert aten ^ export_raw_ir
@@ -294,13 +294,16 @@ def _export_to_pretty_string(model, args, f, export_params=True, verbose=False,
                              input_names=None, output_names=None, operator_export_type=OperatorExportTypes.ONNX,
                              export_type=ExportTypes.PROTOBUF_FILE, example_outputs=None, propagate=False,
                              google_printer=False, opset_version=None, _retain_param_name=False,
-                             do_constant_folding=False, keep_initializers_as_inputs=True, fixed_batch_size=False):
+                             do_constant_folding=False, keep_initializers_as_inputs=None, fixed_batch_size=False):
     from torch.onnx.symbolic_helper import _default_onnx_opset_version, _set_opset_version
     from torch.onnx.symbolic_helper import _set_operator_export_type
     if opset_version is None:
         opset_version = _default_onnx_opset_version
     _set_opset_version(opset_version)
     _set_operator_export_type(operator_export_type)
+    val_keep_init_as_ip = True if keep_initializers_as_inputs is None else keep_initializers_as_inputs
+    if keep_initializers_as_inputs is None and operator_export_type is OperatorExportTypes.ONNX:
+        val_keep_init_as_ip = False
     graph, params_dict, torch_out = _model_to_graph(model, args, verbose,
                                                     training, input_names,
                                                     output_names, operator_export_type,
@@ -309,7 +312,7 @@ def _export_to_pretty_string(model, args, f, export_params=True, verbose=False,
 
     return graph._pretty_print_onnx(params_dict, opset_version, False,
                                     operator_export_type, google_printer,
-                                    keep_initializers_as_inputs)
+                                    val_keep_init_as_ip)
 
 
 # NOTE: the output `torch_out` will contain the output tensors resulting from
@@ -320,7 +323,7 @@ def _export(model, args, f, export_params=True, verbose=False, training=False,
             input_names=None, output_names=None, operator_export_type=OperatorExportTypes.ONNX,
             export_type=ExportTypes.PROTOBUF_FILE, example_outputs=None, propagate=False,
             opset_version=None, _retain_param_name=False, do_constant_folding=False,
-            strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=True, fixed_batch_size=False):
+            strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None, fixed_batch_size=False):
     if isinstance(model, torch.nn.DataParallel):
         raise ValueError('torch.nn.DataParallel is not supported by ONNX '
                          'exporter, please use \'attribute\' module to '
@@ -336,6 +339,9 @@ def _export(model, args, f, export_params=True, verbose=False, training=False,
             opset_version = _default_onnx_opset_version
         _set_opset_version(opset_version)
         _set_operator_export_type(operator_export_type)
+        val_keep_init_as_ip = True if keep_initializers_as_inputs is None else keep_initializers_as_inputs
+        if keep_initializers_as_inputs is None and operator_export_type is OperatorExportTypes.ONNX:
+            val_keep_init_as_ip = False
         graph, params_dict, torch_out = _model_to_graph(model, args, verbose,
                                                         training, input_names,
                                                         output_names, operator_export_type,
@@ -353,11 +359,11 @@ def _export(model, args, f, export_params=True, verbose=False, training=False,
         if export_params:
             proto, export_map = graph._export_onnx(
                 params_dict, opset_version, dynamic_axes, defer_weight_export,
-                operator_export_type, strip_doc_string, keep_initializers_as_inputs)
+                operator_export_type, strip_doc_string, val_keep_init_as_ip)
         else:
             proto, export_map = graph._export_onnx(
                 {}, opset_version, dynamic_axes, False, operator_export_type,
-                strip_doc_string, keep_initializers_as_inputs)
+                strip_doc_string, val_keep_init_as_ip)
 
         if export_type == ExportTypes.PROTOBUF_FILE:
             assert(len(export_map) == 0)