Add a new api from_tflite to improve user experience. (#1954)

* Add a new api from_tflite to improve user experience.

Signed-off-by: Jay Zhang <[email protected]>

Author: fatcat-z

README.md

@@ -377,6 +377,36 @@ model_proto, external_tensor_storage = tf2onnx.convert.from_graph_def(graph_def,
         An ONNX model_proto and an external_tensor_storage dict.
 ```
 
+### from_tflite
+```
+import tf2onnx
+
+model_proto, external_tensor_storage = tf2onnx.convert.from_tflite(tflite_path,
+                input_names=None, output_names=None, opset=None, custom_ops=None, custom_op_handlers=None,
+                custom_rewriter=None, inputs_as_nchw=None, extra_opset=None, shape_override=None, target=None,
+                large_model=False, output_path=None):
+
+    Args:
+        tflite_path: the tflite model file full path
+        input_names: list of input names
+        output_names: list of output names
+        opset: the opset to be used for the ONNX model, default is the latest
+        custom_ops: if a model contains ops not recognized by onnx runtime,
+            you can tag these ops with a custom op domain so that the
+            runtime can still open the model. Type is a dictionary `{op name: domain}`.
+        custom_op_handlers: dictionary of custom ops handlers
+        custom_rewriter: list of custom graph rewriters
+        inputs_as_nchw: transpose inputs in list from nchw to nhwc
+        extra_opset: list of extra opset's, for example the opset's used by custom ops
+        shape_override: dict with inputs that override the shapes given by tensorflow
+        target: list of workarounds applied to help certain platforms
+        large_model: use the ONNX external tensor storage format
+        output_path: save model to output_path
+
+    Returns:
+        An ONNX model_proto and an external_tensor_storage dict.
+```
+
 ### Creating custom op mappings from python
 
 For complex custom ops that require graph rewrites or input / attribute rewrites using the python interface to insert a custom op will be the easiest way to accomplish the task.

tests/models/regression/tflite/test_api_model.tflite

@@ -231,6 +231,19 @@ def test_graphdef(self):
         self.assertTrue(output_names[0] == "pred")
         self.assertAllClose([2.1193342], oy[0], rtol=0.1, atol=0.1)
 
+    @check_tf_min_version("2.0")
+    def test_tflite(self):
+        output_path = os.path.join(self.test_data_directory, "model.onnx")
+
+        x_val = np.array([1.0, 2.0, -3.0, -4.0], dtype=np.float32).reshape((2, 2))
+        model_proto, _ = tf2onnx.convert.from_tflite("tests/models/regression/tflite/test_api_model.tflite",
+                                                     input_names=['input'], output_names=['output'],
+                                                     output_path=output_path)
+        output_names = [n.name for n in model_proto.graph.output]
+        oy = self.run_onnxruntime(output_path, {"input": x_val}, output_names)
+        self.assertTrue(output_names[0] == "output")
+        exp_result = tf.add(x_val, x_val)
+        self.assertAllClose(exp_result, oy[0], rtol=0.1, atol=0.1)
 
 if __name__ == '__main__':
     unittest_main()

tests/test_api.py

@@ -636,5 +636,61 @@ def from_graph_def(graph_def, name=None, input_names=None, output_names=None, op
     return model_proto, external_tensor_storage
 
 
+def from_tflite(tflite_path, input_names=None, output_names=None, opset=None, custom_ops=None, custom_op_handlers=None,
+                custom_rewriter=None, inputs_as_nchw=None, extra_opset=None, shape_override=None, target=None,
+                large_model=False, output_path=None):
+    """Returns a ONNX model_proto for a tflite model file.
+
+    Args:
+        tflite_path: the tflite model file full path
+        input_names: list of input names
+        output_names: list of output names
+        opset: the opset to be used for the ONNX model, default is the latest
+        custom_ops: if a model contains ops not recognized by onnx runtime,
+            you can tag these ops with a custom op domain so that the
+            runtime can still open the model. Type is a dictionary `{op name: domain}`.
+        custom_op_handlers: dictionary of custom ops handlers
+        custom_rewriter: list of custom graph rewriters
+        inputs_as_nchw: transpose inputs in list from nchw to nhwc
+        extra_opset: list of extra opset's, for example the opset's used by custom ops
+        shape_override: dict with inputs that override the shapes given by tensorflow
+        target: list of workarounds applied to help certain platforms
+        large_model: use the ONNX external tensor storage format
+        output_path: save model to output_path
+
+    Returns:
+        An ONNX model_proto and an external_tensor_storage dict.
+    """
+    if not tflite_path:
+        raise ValueError("tflite_path needs to be provided")
+    if not input_names:
+        input_names = []
+    if not output_names:
+        output_names = []
+
+    with tf.device("/cpu:0"):
+        model_proto, external_tensor_storage = _convert_common(
+            None,
+            tflite_path=tflite_path,
+            name=os.path.splitext(os.path.basename(tflite_path))[0],
+            continue_on_error=True,
+            target=target,
+            opset=opset,
+            custom_ops=custom_ops,
+            custom_op_handlers=custom_op_handlers,
+            custom_rewriter=custom_rewriter,
+            extra_opset=extra_opset,
+            shape_override=shape_override,
+            input_names=input_names,
+            output_names=output_names,
+            inputs_as_nchw=inputs_as_nchw,
+            large_model=large_model,
+            tensors_to_rename=None,
+            initialized_tables=None,
+            output_path=output_path)
+
+        return model_proto, external_tensor_storage
+
+
 if __name__ == "__main__":
     main()