Use convert_attribute=True by default for onnx.save (#1738)

* use convert_attribute=True by default for onnx.save * leave some constants in the model.onnx
huggingface · Mar 1, 2024 · ebb63f8 · ebb63f8
1 parent 27deaf6
commit ebb63f8
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 4 deletions.
diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py
@@ -342,7 +342,11 @@ def fix_dynamic_axes(
                 dims = onnx_model.graph.output[output_idx].type.tensor_type.shape.dim
                 dims[dim_idx].dim_value = outputs[output_idx].shape[dim_idx]
 
-            onnx.save(onnx_model, model_path.as_posix())
+            onnx.save(
+                onnx_model,
+                model_path.as_posix(),
+                convert_attribute=True,
+            )
             del onnx_model
             gc.collect()
 

diff --git a/optimum/exporters/onnx/convert.py b/optimum/exporters/onnx/convert.py
@@ -610,7 +610,8 @@ def remap(value):
                 save_as_external_data=True,
                 all_tensors_to_one_file=True,
                 location=output.name + "_data",
-                size_threshold=1024 if not FORCE_ONNX_EXTERNAL_DATA else 0,
+                size_threshold=1024 if not FORCE_ONNX_EXTERNAL_DATA else 100,
+                convert_attribute=True,
             )
 
             # delete previous external data
@@ -688,7 +689,11 @@ def export_tensorflow(
 
     with config.patch_model_for_export(model):
         onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature, opset=opset)
-    onnx.save(onnx_model, output.as_posix())
+    onnx.save(
+        onnx_model,
+        output.as_posix(),
+        convert_attribute=True,
+    )
 
     return input_names, output_names
 

diff --git a/optimum/onnx/graph_transformations.py b/optimum/onnx/graph_transformations.py
@@ -157,7 +157,11 @@ def check_and_save_model(model: onnx.ModelProto, save_path: Optional[Union[str,
                 # The new model may be below the maximum protobuf size, overwritting a model that was larger. Hence this os.remove.
                 os.remove(external_path)
 
-            onnx.save(model, save_path)
+            onnx.save(
+                model,
+                save_path,
+                convert_attribute=True,
+            )
     elif save_path is not None:
         # path/to/model.onnx
         save_path = Path(save_path).as_posix()
@@ -177,6 +181,7 @@ def check_and_save_model(model: onnx.ModelProto, save_path: Optional[Union[str,
             save_as_external_data=True,
             all_tensors_to_one_file=True,
             location=external_file_name,
+            convert_attribute=True,
         )
         try:
             onnx.checker.check_model(save_path)

diff --git a/optimum/onnxruntime/modeling_decoder.py b/optimum/onnxruntime/modeling_decoder.py
@@ -536,6 +536,7 @@ def _from_pretrained(
                 all_tensors_to_one_file=True,
                 location=model_cache_path.name + "_data",
                 size_threshold=0,
+                convert_attribute=True,
             )
         del onnx_model
 

diff --git a/tests/onnx/test_onnx_graph_transformations.py b/tests/onnx/test_onnx_graph_transformations.py
@@ -117,6 +117,7 @@ def test_to_int32(self):
                 save_as_external_data=True,
                 all_tensors_to_one_file=True,
                 location=Path(save_path).name + "_data",
+                convert_attribute=True,
             )
 
             onnx.checker.check_model(save_path)