Add helper function update_inputs_outputs_dims to tools (#2148)

BowenBao · wschin · commit 599f3da985cf · 2019-09-13T15:20:29.000-07:00
* Add a helper function update_inputs_outputs_dims to tools

* fix link to doc

* newline at the end

* add test for tools

* doc props

* nit

* ci tests

* ci tests 2

* accept shapes by dictionary inputs and add more error handling

* Update onnx/tools/update_model_dims.py

nit: rephrasing

Co-Authored-By: Wei-Sheng Chin &lt;wschin@outlook.com&gt;

* remove debug line

* fix type annotation

* fix annotation

* fix annotation

* fix annotation

* fix flake8
diff --git a/docs/PythonAPIOverview.md b/docs/PythonAPIOverview.md
@@ -214,3 +214,19 @@ import onnx.utils
 model = onnx.load('path/to/the/model.onnx')
 polished_model = onnx.utils.polish_model(model)
 ```
+
+## Tools
+### Updating Model's Inputs Outputs Dimension Sizes with Variable Length
+Function `update_inputs_outputs_dims` updates the dimension of the inputs and outputs of the model,
+to the provided values in the parameter. You could provide both static and dynamic dimension size,
+by using dim_param. For more information on static and dynamic dimension size, checkout [Tensor Shapes](IR.md#tensor-shapes).
+
+The function runs model checker after the input/output sizes are updated.
+```python
+import onnx
+from onnx.tools import update_model_dims
+
+model = onnx.load('path/to/the/model.onnx')
+# Here both 'seq', 'batch' and -1 are dynamic using dim_param.
+variable_length_model = update_model_dims.update_inputs_outputs_dims(model, {'input_name': ['seq', 'batch', 3, -1]}, {'output_name': ['seq', 'batch', 1, -1]})
+```
diff --git a/onnx/test/tools_test.py b/onnx/test/tools_test.py
@@ -0,0 +1,46 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import unittest
+import onnx
+from onnx.tools import update_model_dims
+from onnx import helper, TensorProto
+
+
+class TestToolsFunctions(unittest.TestCase):
+    def test_update_inputs_outputs_dim(self):  # type: () -> None
+        node_def = helper.make_node(
+            "Conv",
+            inputs=['x', 'W'],
+            outputs=['y'],
+            kernel_shape=[3, 3],
+            strides=[2, 2],
+        )
+        graph_def = helper.make_graph(
+            [node_def],
+            'test',
+            [helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 1, 5, 5]),
+             helper.make_tensor_value_info('W', TensorProto.FLOAT, [1, 1, 3, 3])],
+            [helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 1, 2, 2])]
+        )
+        model_def = helper.make_model(graph_def, producer_name='test')
+        updated_def = update_model_dims.update_inputs_outputs_dims(
+            model_def,
+            {
+                "x": [1, 1, 'x1', -1],
+                "W": [1, 1, 3, 3],
+            },
+            {
+                "y": [1, 1, -1, -1],
+            })
+        onnx.checker.check_model(updated_def)
+        self.assertEqual(updated_def.graph.input[0].type.tensor_type.shape.dim[2].dim_param, 'x1')
+        self.assertEqual(updated_def.graph.input[0].type.tensor_type.shape.dim[3].dim_param, 'x_3')
+        self.assertEqual(updated_def.graph.output[0].type.tensor_type.shape.dim[2].dim_param, 'y_2')
+        self.assertEqual(updated_def.graph.output[0].type.tensor_type.shape.dim[3].dim_param, 'y_3')
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/onnx/tools/update_model_dims.py b/onnx/tools/update_model_dims.py
@@ -0,0 +1,83 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from six import string_types
+from typing import Any, List, Text, Dict, Set
+from onnx import ModelProto, ValueInfoProto
+
+import onnx.checker
+
+
+def update_inputs_outputs_dims(model, input_dims, output_dims):  # type: (ModelProto, Dict[Text, List[Any]], Dict[Text, List[Any]]) -> ModelProto
+    """
+        This function updates the dimension sizes of the model's inputs and outputs to the values
+        provided in input_dims and output_dims. if the dim value provided is negative, a unique dim_param
+        will be set for that dimension.
+
+        Example. if we have the following shape for inputs and outputs:
+                shape(input_1) = ('b', 3, 'w', 'h')
+                shape(input_2) = ('b', 4)
+                and shape(output)  = ('b', 'd', 5)
+
+            The parameters can be provided as:
+                input_dims = {
+                    "input_1": ['b', 3, 'w', 'h'],
+                    "input_2": ['b', 4],
+                }
+                output_dims = {
+                    "output": ['b', -1, 5]
+                }
+
+            Putting it together:
+                model = onnx.load('model.onnx')
+                updated_model = update_inputs_outputs_dims(model, input_dims, output_dims)
+                onnx.save(updated_model, 'model.onnx')
+    """
+    dim_param_set = set()  # type: Set[Text]
+
+    def init_dim_param_set(dim_param_set, value_infos):  # type: (Set[Text], List[ValueInfoProto]) -> None
+        for info in value_infos:
+            shape = info.type.tensor_type.shape
+            for dim in shape.dim:
+                if dim.HasField('dim_param'):
+                    dim_param_set.add(dim.dim_param)  # type: ignore
+
+    init_dim_param_set(dim_param_set, model.graph.input)  # type: ignore
+    init_dim_param_set(dim_param_set, model.graph.output)  # type: ignore
+    init_dim_param_set(dim_param_set, model.graph.value_info)  # type: ignore
+
+    def update_dim(tensor, dim, j, name):  # type: (ValueInfoProto, Any, int, Text) -> None
+        dim_proto = tensor.type.tensor_type.shape.dim[j]
+        if isinstance(dim, int):
+            if dim >= 0:
+                if dim_proto.HasField('dim_value') and dim_proto.dim_value != dim:
+                    raise ValueError('Unable to set dimension value to {} for axis {} of {}. Contradicts existing dimension value {}.'
+                        .format(dim, j, name, dim_proto.dim_value))
+                dim_proto.dim_value = dim
+            else:
+                generated_dim_param = name + '_' + str(j)
+                if generated_dim_param in dim_param_set:
+                    raise ValueError('Unable to generate unique dim_param for axis {} of {}. Please manually provide a dim_param value.'
+                        .format(j, name))
+                dim_proto.dim_param = generated_dim_param
+        elif isinstance(dim, string_types):
+            dim_proto.dim_param = dim
+        else:
+            raise ValueError('Only int or str is accepted as dimension value, incorrect type: {}'.format(type(dim)))
+
+    for input in model.graph.input:
+        input_name = input.name
+        input_dim_arr = input_dims[input_name]
+        for j, dim in enumerate(input_dim_arr):
+            update_dim(input, dim, j, input_name)
+
+    for output in model.graph.output:
+        output_name = output.name
+        output_dim_arr = output_dims[output_name]
+        for j, dim in enumerate(output_dim_arr):
+            update_dim(output, dim, j, output_name)
+
+    onnx.checker.check_model(model)
+    return model
