Repository: onnx/onnx
Pull request: Modify existing function in external data helper
Signed-off-by: Anna Jung (VMware) <antheaj@vmware.com>

Update external data helpers to modify an existing function `convert_model_to_external_data` to take a `model_file_path` and `convert_attribute` to convert attributes to external data only when specified to `true` and to save the model when `model_file_path` is specified

fixes #3228 

assign @askhade for review

Problem statement:
Updates to external data helpers
Add a new API for converting a model to external data. Today the conversion happens in 2 steps
`external_data_helper.convert_model_to_external_data(<model>,
                                                        <all_tensors_to_one_file>, <size_threshold>)
 save_model(model, output_path)
`
We want to add another api which combines the 2 steps
`
save_model_to_external_data(<model>, <output_path>, <all_tensors_to_one_file>, <size_threshold>, <convert_attributes>)
`

Hints from issue discussion:
Hi @askhade, thanks for suggesting me to work on this issue! 
I have a question about `conver_attributes` parameter in the `save_model_to_external_data`.

Neither of `convert_model_to_external_data` and `save_model` has this input. So wondering what your thoughts are in including this in the combined API `save_model_to_external_data`. 

Could you help expand the input and how `convert_attributes` should be used?
Sure!
If you look at the external data helpers code you will see we convert every tesnsor to external data even the ones in AttributeProto. This is not ideal since attributes are not usually large and it takes more memory to make them external than just leaving then in the model itself. So by default we should to exclude tensorprotos in attributeproto but also give an option to convert them to external data to preserve existing behaviour. Therefore I am proposing to add a boolean convert_attributes. 
Thanks for the explanation! I will take a look and let you know if I have more questions
Hi @askhade I want to run the logic by you and see if this makes sense.

The `save_model_to_external_data` function takes parameters listed above. 

When the parameter `convert_attribute` is true, it runs the existing `convert_model_to_external_data` function then saves the model. 
If false, it runs the similar code but tensors loops through `_get_initializer_tensors` function instead `_get_all_tensors` function then saves the model. 

Does that logic seem okay to you? I wanted to double-check because I'm running into a few problems with unit tests that I am writing for the `save_model_to_external_data`. There seems to be no difference between the test result when using  `_get_initializer_tensors` vs `_get_all_tensors`. 

In case my logic is wrong, I wanted to get your feedback to determine where to focus on. Any feedback would be appreciated! thanks :D 
Update: I was a little confused, but found my way! I got confused on why some had `data_location` and some did not. Regardless, only the converted has the field which I can use for testing. I ended up using `MessageToJson`, but not sure if this is the best way. Will try to wrap up and make a PR soon for you to take a look at.
> Hi @askhade I want to run the logic by you and see if this makes sense.
> 
> The `save_model_to_external_data` function takes parameters listed above.
> 
> When the parameter `convert_attribute` is true, it runs the existing `convert_model_to_external_data` function then saves the model.
> If false, it runs the similar code but tensors loops through `_get_initializer_tensors` function instead `_get_all_tensors` function then saves the model.
> 
> Does that logic seem okay to you? I wanted to double-check because I'm running into a few problems with unit tests that I am writing for the `save_model_to_external_data`. There seems to be no difference between the test result when using `_get_initializer_tensors` vs `_get_all_tensors`.
> 
> In case my logic is wrong, I wanted to get your feedback to determine where to focus on. Any feedback would be appreciated! thanks :D

for you to see the difference you need to make sure your test case includes an attribute which has data stored in tensorproto otherwise there is no difference. For example if you can test a model with ConstantOfShape node you will see the difference when convert_attribute is set to true vs false.

I realize that it's not the operator that was the problem and kept it simple and check for the "data_location" field, hopefully, that's good validation. PR is open to get your feedback :) 

PR: https://github.com/onnx/onnx/pull/3280

Commit message: Modify existing save model api to optionally convert and save as external data

Signed-off-by: Anna Jung (VMware) <antheaj@vmware.com>

Patch to review:
diff --git a/docs/ExternalData.md b/docs/ExternalData.md
index 7d35ac8b84f..c1e53d7bef6 100644
--- a/docs/ExternalData.md
+++ b/docs/ExternalData.md
@@ -22,6 +22,27 @@ load_external_data_for_model(onnx_model, 'data/directory/path/')
 # Then the onnx_model has loaded the external data from the specific directory
 ```
 
+## Converting an ONNX Model to External Data
+```python
+import onnx
+from onnx.external_data_helper import convert_model_to_external_data
+
+onnx_model = ... # Your model in memory as ModelProto
+convert_model_to_external_data(onnx_model, all_tensors_to_one_file=True, location='filename', size_threshold=1024, convert_attribute=False)
+# Must be followed by save_model to save the converted model to a specific path
+onnx.save_model(onnx_model, 'path/to/save/the/model.onnx')
+# Then the onnx_model has converted raw data as external data and saved to specific directory
+```
+
+## Converting and Saving an ONNX Model to External Data
+```python
+import onnx
+
+onnx_model = ... # Your model in memory as ModelProto
+onnx.save_model(onnx_model, 'path/to/save/the/model.onnx', save_as_external_data=True, all_tensors_to_one_file=True, location='filename', size_threshold=1024, convert_attribute=False)
+# Then the onnx_model has converted raw data as external data and saved to specific directory
+```
+
 ## onnx.checker for Models with External Data
 
 ### Models with External Data (<2GB)
diff --git a/docs/PythonAPIOverview.md b/docs/PythonAPIOverview.md
index 34b7eaa3846..db9d1c6cc02 100644
--- a/docs/PythonAPIOverview.md
+++ b/docs/PythonAPIOverview.md
@@ -33,6 +33,16 @@ load_external_data_for_model(onnx_model, 'data/directory/path/')
 # Then the onnx_model has loaded the external data from the specific directory
 ```
 
+## Converting an ONNX Model to External Data
+```python
+from onnx.external_data_helper import convert_model_to_external_data
+
+onnx_model = ... # Your model in memory as ModelProto
+convert_model_to_external_data(onnx_model, all_tensors_to_one_file=True, location='filename', size_threshold=1024, convert_attribute=False)
+# Then the onnx_model has converted raw data as external data
+# Must be followed by save
+```
+
 ## Saving an ONNX Model
 ```python
 import onnx
@@ -45,6 +55,17 @@ onnx.save(onnx_model, 'path/to/the/model.onnx')
 Runnable IPython notebooks:
 - [save_model.ipynb](https://github.com/onnx/onnx/tree/master/onnx/examples/save_model.ipynb)
 
+
+## Converting and Saving an ONNX Model to External Data
+```python
+import onnx
+
+onnx_model = ... # Your model in memory as ModelProto
+onnx.save_model(onnx_model, 'path/to/save/the/model.onnx', save_as_external_data=True, all_tensors_to_one_file=True, location='filename', size_threshold=1024, convert_attribute=False)
+# Then the onnx_model has converted raw data as external data and saved to specific directory
+```
+
+
 ## Manipulating TensorProto and Numpy Array
 ```python
 import numpy
diff --git a/onnx/__init__.py b/onnx/__init__.py
index 439853590d1..5bd0a80fa73 100644
--- a/onnx/__init__.py
+++ b/onnx/__init__.py
@@ -8,7 +8,7 @@
 import os
 
 from .onnx_cpp2py_export import ONNX_ML
-from onnx.external_data_helper import load_external_data_for_model, write_external_data_tensors
+from onnx.external_data_helper import load_external_data_for_model, write_external_data_tensors, convert_model_to_external_data
 from .onnx_pb import *  # noqa
 from .onnx_operators_pb import * # noqa
 from .onnx_data_pb import * # noqa
@@ -172,18 +172,29 @@ def load_tensor_from_string(s, format=None):  # type: (bytes, Optional[Any]) ->
     return _deserialize(s, TensorProto())
 
 
-def save_model(proto, f, format=None):  # type: (Union[ModelProto, bytes], Union[IO[bytes], Text], Optional[Any]) -> None
+def save_model(proto, f, format=None, save_as_external_data=False, all_tensors_to_one_file=True, location=None, size_threshold=1024, convert_attribute=False):
+    # type: (Union[ModelProto, bytes], Union[IO[bytes], Text], Optional[Any], bool, bool, Optional[Text], int, bool) -> None
     '''
-    Saves the ModelProto to the specified path.
+    Saves the ModelProto to the specified path and optionally, set tensors with raw data as external data before saving.
 
     @params
-    proto should be a in-memory ModelProto
-    f can be a file-like object (has "write" function) or a string containing a file name
-    format is for future use
+    proto: should be a in-memory ModelProto
+    f: can be a file-like object (has "write" function) or a string containing a file name format is for future use
+    all_tensors_to_one_file: If true, save all tensors to one external file specified by location.
+                             If false, save each tensor to a file named with the tensor name.
+    location: specify the external file that all tensors to save to.
+              If not specified, will use the model name.
+    size_threshold: Threshold for size of data. Only when tensor's data is >= the size_threshold it will be converted
+                    to external data. To convert every tensor with raw data to external data set size_threshold=0.
+    convert_attribute: If true, convert all tensors to external data
+                       If false, convert only non-attribute tensors to external data
     '''
     if isinstance(proto, bytes):
         proto = _deserialize(proto, ModelProto())
 
+    if save_as_external_data:
+        convert_model_to_external_data(proto, all_tensors_to_one_file, location, size_threshold, convert_attribute)
+
     model_filepath = _get_file_path(f)
     if model_filepath:
         basepath = os.path.dirname(model_filepath)
diff --git a/onnx/external_data_helper.py b/onnx/external_data_helper.py
index 9bf6f1866c5..86d235d86a8 100644
--- a/onnx/external_data_helper.py
+++ b/onnx/external_data_helper.py
@@ -10,6 +10,7 @@
 import sys
 from itertools import chain
 from typing import Iterable, Text, Optional
+
 from .onnx_pb import TensorProto, ModelProto
 
 
@@ -99,8 +100,8 @@ def set_external_data(tensor,  # type: TensorProto
             entry.value = str(v)
 
 
-def convert_model_to_external_data(model, all_tensors_to_one_file=True, location=None, size_threshold=1024):
-    # type: (ModelProto, bool, Optional[Text], int) -> None
+def convert_model_to_external_data(model, all_tensors_to_one_file=True, location=None, size_threshold=1024, convert_attribute=False):
+    # type: (ModelProto, bool, Optional[Text], int, bool) -> None
     """
     Call to set all tensors with raw data as external data. This call should preceed 'save_model'.
     'save_model' saves all the tensors data as external data after calling this function.
@@ -112,16 +113,22 @@ def convert_model_to_external_data(model, all_tensors_to_one_file=True, location
               If not specified, will use the model name.
     size_threshold: Threshold for size of data. Only when tensor's data is >= the size_threshold
     it will be converted to external data. To convert every tensor with raw data to external data set size_threshold=0.
+    convert_attribute: If true, convert all tensors to external data
+                       If false, convert only non-attribute tensors to external data
     """
+    tensors = _get_initializer_tensors(model)
+    if convert_attribute:
+        tensors = _get_all_tensors(model)
+
     if all_tensors_to_one_file:
         file_name = Text(uuid.uuid1())
         if location:
             file_name = location
-        for tensor in _get_all_tensors(model):
+        for tensor in tensors:
             if tensor.HasField("raw_data") and sys.getsizeof(tensor.raw_data) >= size_threshold:
                 set_external_data(tensor, file_name)
     else:
-        for tensor in _get_all_tensors(model):
+        for tensor in tensors:
             if tensor.HasField("raw_data") and sys.getsizeof(tensor.raw_data) >= size_threshold:
                 tensor_location = tensor.name
                 if not _is_valid_filename(tensor_location):
diff --git a/onnx/test/test_external_data.py b/onnx/test/test_external_data.py
index a0bf86ce57d..428b5abb887 100644
--- a/onnx/test/test_external_data.py
+++ b/onnx/test/test_external_data.py
@@ -203,19 +203,36 @@ def create_test_model_proto(self):  # type: () -> ModelProto
     def test_check_model(self):  # type: () -> None
         checker.check_model(self.model)
 
-    def test_convert_model_to_from_one_file(self):  # type: () -> None
+    def test_convert_model_to_external_data_with_size_threshold(self):  # type: () -> None
         model_file_path = self.get_temp_model_filename()
-        external_data_file = str(uuid.uuid4())
-        convert_model_to_external_data(self.model, location=external_data_file, size_threshold=0)
+
+        convert_model_to_external_data(self.model, size_threshold=1024)
         onnx.save_model(self.model, model_file_path)
-        self.assertTrue(Path.isfile(model_file_path))
-        self.assertTrue(Path.isfile(os.path.join(self.temp_dir, external_data_file)))
+
+        model = onnx.load_model(model_file_path)
+        initializer_tensor = model.graph.initializer[0]
+        self.assertFalse(initializer_tensor.HasField("data_location"))
+
+    def test_convert_model_to_external_data_without_size_threshold(self):  # type: () -> None
+        model_file_path = self.get_temp_model_filename()
+        convert_model_to_external_data(self.model, size_threshold=0)
+        onnx.save_model(self.model, model_file_path)
+
         model = onnx.load_model(model_file_path)
         initializer_tensor = model.graph.initializer[0]
+        self.assertTrue(initializer_tensor.HasField("data_location"))
         self.assertTrue(np.allclose(to_array(initializer_tensor), self.initializer_value))
 
-        attribute_tensor = model.graph.node[0].attribute[0].t
-        self.assertTrue(np.allclose(to_array(attribute_tensor), self.attribute_value))
+    def test_convert_model_to_external_data_from_one_file_with_location(self):  # type: () -> None
+        model_file_path = self.get_temp_model_filename()
+        external_data_file = str(uuid.uuid4())
+
+        convert_model_to_external_data(self.model, size_threshold=0, all_tensors_to_one_file=True, location=external_data_file)
+        onnx.save_model(self.model, model_file_path)
+
+        self.assertTrue(Path.isfile(os.path.join(self.temp_dir, external_data_file)))
+
+        model = onnx.load_model(model_file_path)
 
         # test convert model from external data
         convert_model_from_external_data(model)
@@ -228,36 +245,101 @@ def test_convert_model_to_from_one_file(self):  # type: () -> None
         self.assertTrue(np.allclose(to_array(initializer_tensor), self.initializer_value))
 
         attribute_tensor = model.graph.node[0].attribute[0].t
-        self.assertFalse(len(initializer_tensor.external_data))
+        self.assertFalse(len(attribute_tensor.external_data))
         self.assertEqual(attribute_tensor.data_location, TensorProto.DEFAULT)
         self.assertTrue(np.allclose(to_array(attribute_tensor), self.attribute_value))
 
-    def test_convert_model_to_external_data_one_file_per_tensor(self):  # type: () -> None
+    def test_convert_model_to_external_data_from_one_file_without_location_uses_model_name(self):  # type: () -> None
         model_file_path = self.get_temp_model_filename()
-        convert_model_to_external_data(self.model, all_tensors_to_one_file=False, size_threshold=0)
+
+        convert_model_to_external_data(self.model, size_threshold=0, all_tensors_to_one_file=True)
         onnx.save_model(self.model, model_file_path)
+
+        self.assertTrue(Path.isfile(model_file_path))
+        self.assertTrue(Path.isfile(os.path.join(self.temp_dir, model_file_path)))
+
+    def test_convert_model_to_external_data_one_file_per_tensor_without_attribute(self):  # type: () -> None
+        model_file_path = self.get_temp_model_filename()
+
+        convert_model_to_external_data(self.model, size_threshold=0, all_tensors_to_one_file=False, convert_attribute=False)
+        onnx.save_model(self.model, model_file_path)
+
+        self.assertTrue(Path.isfile(model_file_path))
+        self.assertTrue(Path.isfile(os.path.join(self.temp_dir, "input_value")))
+        self.assertFalse(Path.isfile(os.path.join(self.temp_dir, "attribute_value")))
+
+    def test_convert_model_to_external_data_one_file_per_tensor_with_attribute(self):  # type: () -> None
+        model_file_path = self.get_temp_model_filename()
+
+        convert_model_to_external_data(self.model, size_threshold=0, all_tensors_to_one_file=False, convert_attribute=True)
+        onnx.save_model(self.model, model_file_path)
+
         self.assertTrue(Path.isfile(model_file_path))
         self.assertTrue(Path.isfile(os.path.join(self.temp_dir, "input_value")))
         self.assertTrue(Path.isfile(os.path.join(self.temp_dir, "attribute_value")))
+
+    def test_convert_model_to_external_data_does_not_convert_attribute_values(self):  # type: () -> None
+        model_file_path = self.get_temp_model_filename()
+
+        convert_model_to_external_data(self.model, size_threshold=0, convert_attribute=False, all_tensors_to_one_file=False)
+        onnx.save_model(self.model, model_file_path)
+
+        self.assertTrue(Path.isfile(os.path.join(self.temp_dir, "input_value")))
+        self.assertFalse(Path.isfile(os.path.join(self.temp_dir, "attribute_value")))
+
+        model = onnx.load_model(model_file_path)
+        initializer_tensor = model.graph.initializer[0]
+        self.assertTrue(initializer_tensor.HasField("data_location"))
+
+        attribute_tensor = model.graph.node[0].attribute[0].t
+        self.assertFalse(attribute_tensor.HasField("data_location"))
+
+    def test_convert_model_to_external_data_converts_attribute_values(self):  # type: () -> None
+        model_file_path = self.get_temp_model_filename()
+
+        convert_model_to_external_data(self.model, size_threshold=0, convert_attribute=True)
+        onnx.save_model(self.model, model_file_path)
+
         model = onnx.load_model(model_file_path)
+
         initializer_tensor = model.graph.initializer[0]
         self.assertTrue(np.allclose(to_array(initializer_tensor), self.initializer_value))
+        self.assertTrue(initializer_tensor.HasField("data_location"))
 
         attribute_tensor = model.graph.node[0].attribute[0].t
         self.assertTrue(np.allclose(to_array(attribute_tensor), self.attribute_value))
+        self.assertTrue(attribute_tensor.HasField("data_location"))
 
-    def test_convert_model_to_external_data_with_size_threshold(self):  # type: () -> None
+    def test_save_model_does_not_convert_to_external_data_and_saves_the_model(self):  # type: () -> None
         model_file_path = self.get_temp_model_filename()
-        convert_model_to_external_data(self.model, all_tensors_to_one_file=False, size_threshold=1024)
-        onnx.save_model(self.model, model_file_path)
+        onnx.save_model(self.model, model_file_path, save_as_external_data=False)
         self.assertTrue(Path.isfile(model_file_path))
-        self.assertFalse(Path.isfile(os.path.join(self.temp_dir, "input_value")))
-        self.assertFalse(Path.isfile(os.path.join(self.temp_dir, "attribute_value")))
+
         model = onnx.load_model(model_file_path)
         initializer_tensor = model.graph.initializer[0]
+        self.assertFalse(initializer_tensor.HasField("data_location"))
+
+        attribute_tensor = model.graph.node[0].attribute[0].t
+        self.assertFalse(attribute_tensor.HasField("data_location"))
+
+    def test_save_model_does_convert_and_saves_the_model(self):  # type: () -> None
+        model_file_path = self.get_temp_model_filename()
+        onnx.save_model(self.model,
+                        model_file_path,
+                        save_as_external_data=True,
+                        all_tensors_to_one_file=True,
+                        location=None,
+                        size_threshold=0,
+                        convert_attribute=False)
+
+        model = onnx.load_model(model_file_path)
+
+        initializer_tensor = model.graph.initializer[0]
+        self.assertTrue(initializer_tensor.HasField("data_location"))
         self.assertTrue(np.allclose(to_array(initializer_tensor), self.initializer_value))
 
         attribute_tensor = model.graph.node[0].attribute[0].t
+        self.assertFalse(attribute_tensor.HasField("data_location"))
         self.assertTrue(np.allclose(to_array(attribute_tensor), self.attribute_value))
 
 


Code review comments:
diff_hunk: @@ -172,18 +172,29 @@ def load_tensor_from_string(s, format=None):  # type: (bytes, Optional[Any]) ->
     return _deserialize(s, TensorProto())
 
 
-def save_model(proto, f, format=None):  # type: (Union[ModelProto, bytes], Union[IO[bytes], Text], Optional[Any]) -> None
+def save_model(proto, f, format=None, save_as_external_data=False, all_tensors_to_one_file=True, location=None, size_threshold=1024, convert_attribute=False):
+    # type: (Union[ModelProto, bytes], Union[IO[bytes], Text], Optional[Any], bool, bool, Optional[Text], int, bool) -> None
     '''
-    Saves the ModelProto to the specified path.
+    Saves the ModelProto to the specified path and optionally, set tensors with raw data as external data before saving.
path: onnx/__init__.py
text: @user1:
nit: set -> serialize