Add support for saved model conversion (#655)

lgeiger · Tombana · web-flow · commit 6bb5a43f11fc · 2021-06-09T10:23:29.000+01:00
* Add support for saved model conversion

Co-Authored-By: Tom Bannink &lt;Tombana@users.noreply.github.com&gt;

* Fix converter tests

* Imrpove target docstrings

* Add missing license comment

Co-authored-by: Tom Bannink &lt;Tombana@users.noreply.github.com&gt;
diff --git a/larq_compute_engine/__init__.py b/larq_compute_engine/__init__.py
@@ -1,4 +1,7 @@
-from larq_compute_engine.mlir.python.converter import convert_keras_model
+from larq_compute_engine.mlir.python.converter import (
+    convert_keras_model,
+    convert_saved_model,
+)
 from larq_compute_engine.tflite.python import interpreter as testing
 
 try:
@@ -9,4 +12,4 @@
 
 __version__ = metadata.version("larq_compute_engine")
 
-__all__ = ["convert_keras_model", "testing"]
+__all__ = ["convert_keras_model", "convert_saved_model", "testing"]
diff --git a/larq_compute_engine/mlir/BUILD b/larq_compute_engine/mlir/BUILD
@@ -335,13 +335,19 @@ tf_cc_binary(
 )
 
 pybind_extension(
-    name = "_graphdef_tfl_flatbuffer",
-    srcs = ["python/graphdef_tfl_flatbuffer.cc"],
-    module_name = "graphdef_tfl_flatbuffer",
+    name = "_tf_tfl_flatbuffer",
+    srcs = [
+        "python/graphdef_tfl_flatbuffer.cc",
+        "python/pybind_export.cc",
+        "python/saved_model_tfl_flatbuffer.cc",
+    ],
+    module_name = "tf_tfl_flatbuffer",
     deps = [
         ":lce_tfl_passes",
         ":tf_to_tfl_flatbuffer",
         "@org_tensorflow//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "@org_tensorflow//tensorflow/compiler/mlir/lite:tf_to_tfl_flatbuffer",
+        "@org_tensorflow//tensorflow/compiler/mlir/lite/python:tf_tfl_flatbuffer_helpers",
         "@org_tensorflow//tensorflow/compiler/mlir/tensorflow:convert_graphdef",
         "@org_tensorflow//tensorflow/compiler/mlir/tensorflow:import_utils",
         "@org_tensorflow//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
@@ -366,7 +372,7 @@ py_library(
         ":tflite_schema_py",
     ],
     deps = [
-        ":_graphdef_tfl_flatbuffer",
+        ":_tf_tfl_flatbuffer",
     ],
 )
 
diff --git a/larq_compute_engine/mlir/python/converter.py b/larq_compute_engine/mlir/python/converter.py
@@ -1,11 +1,15 @@
+import os
 from packaging import version
 import warnings
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 import tensorflow as tf
+import tempfile
 
-from larq_compute_engine.mlir._graphdef_tfl_flatbuffer import (
+from larq_compute_engine.mlir._tf_tfl_flatbuffer import (
     convert_graphdef_to_tflite_flatbuffer,
+    convert_saved_model_to_tflite_flatbuffer,
 )
+
 from larq_compute_engine.mlir.python.util import modify_integer_quantized_model_io_type
 
 from tensorflow.core.framework.types_pb2 import DataType
@@ -57,6 +61,120 @@ def _contains_training_quant_op(graph_def):
     return False
 
 
+def _validate_options(
+    *,
+    inference_input_type=None,
+    inference_output_type=None,
+    target=None,
+    experimental_default_int8_range=None,
+):
+    if inference_input_type not in (tf.float32, tf.int8):
+        raise ValueError(
+            "Expected `inference_input_type` to be either `tf.float32` or `tf.int8`, "
+            f"got {inference_input_type}."
+        )
+    if inference_output_type not in (tf.float32, tf.int8):
+        raise ValueError(
+            "Expected `inference_output_type` to be either `tf.float32` or `tf.int8`, "
+            f"got {inference_output_type}."
+        )
+    if target not in ("arm", "xcore"):
+        raise ValueError(f'Expected `target` to be "arm" or "xcore", but got {target}.')
+
+    if not tf.executing_eagerly():
+        raise RuntimeError(
+            "Graph mode is not supported. Please enable eager execution using "
+            "tf.enable_eager_execution() when using TensorFlow 1.x"
+        )
+    if experimental_default_int8_range:
+        warnings.warn(
+            "Using `experimental_default_int8_range` as fallback quantization stats. "
+            "This should only be used for latency tests."
+        )
+
+
+def convert_saved_model(
+    saved_model_dir: Union[str, os.PathLike],
+    *,  # Require remaining arguments to be keyword-only.
+    inference_input_type: tf.DType = tf.float32,
+    inference_output_type: tf.DType = tf.float32,
+    target: str = "arm",
+    experimental_default_int8_range: Optional[Tuple[float, float]] = None,
+    experimental_enable_bitpacked_activations: bool = False,
+) -> bytes:
+    """Converts a SavedModel to TFLite flatbuffer.
+
+    !!! example
+        ```python
+        tflite_model = convert_saved_model(saved_model_dir)
+        with open("/tmp/my_model.tflite", "wb") as f:
+            f.write(tflite_model)
+        ```
+
+    # Arguments
+        saved_model_dir: SavedModel directory to convert.
+        inference_input_type: Data type of the input layer. Defaults to `tf.float32`,
+            must be either `tf.float32` or `tf.int8`.
+        inference_output_type: Data type of the output layer. Defaults to `tf.float32`,
+            must be either `tf.float32` or `tf.int8`.
+        target: Target hardware platform. Defaults to "arm", must be either "arm"
+            or "xcore".
+        experimental_default_int8_range: Tuple of integers representing `(min, max)`
+            range values for all arrays without a specified range. Intended for
+            experimenting with quantization via "dummy quantization". (default None)
+        experimental_enable_bitpacked_activations: Enable an experimental
+            converter optimisation that attempts to reduce intermediate
+            activation memory usage by bitpacking the activation tensor between
+            consecutive binary convolutions where possible.
+
+    # Returns
+        The converted data in serialized format.
+    """
+    if version.parse(tf.__version__) < version.parse("2.2"):
+        raise RuntimeError(
+            "TensorFlow 2.2 or newer is required for saved model conversion."
+        )
+    _validate_options(
+        inference_input_type=inference_input_type,
+        inference_output_type=inference_output_type,
+        target=target,
+        experimental_default_int8_range=experimental_default_int8_range,
+    )
+
+    saved_model_dir = str(saved_model_dir)
+    saved_model_tags = [tf.saved_model.SERVING]
+    saved_model_exported_names = [tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
+
+    from tensorflow.python.saved_model import loader_impl
+
+    saved_model_pb, _ = loader_impl.parse_saved_model_with_debug_info(saved_model_dir)
+
+    saved_model_version = saved_model_pb.saved_model_schema_version
+    if saved_model_version not in (1, 2):
+        raise ValueError(
+            f"SavedModel file format({saved_model_version}) is not supported"
+        )
+
+    tflite_buffer = convert_saved_model_to_tflite_flatbuffer(
+        saved_model_dir,
+        saved_model_tags,
+        saved_model_exported_names,
+        saved_model_version,
+        target,
+        experimental_default_int8_range,
+        experimental_enable_bitpacked_activations,
+    )
+
+    if inference_input_type != tf.float32 or inference_output_type != tf.float32:
+        tflite_buffer = modify_integer_quantized_model_io_type(
+            tflite_buffer,
+            inference_input_type=inference_input_type,
+            inference_output_type=inference_output_type,
+        )
+
+    return tflite_buffer
+
+
 def convert_keras_model(
     model: tf.keras.Model,
     *,  # Require remaining arguments to be keyword-only.
@@ -81,7 +199,8 @@ def convert_keras_model(
             must be either `tf.float32` or `tf.int8`.
         inference_output_type: Data type of the output layer. Defaults to `tf.float32`,
             must be either `tf.float32` or `tf.int8`.
-        target: Target hardware platform. Must be "arm" or "xcore".
+        target: Target hardware platform. Defaults to "arm", must be either "arm"
+            or "xcore".
         experimental_default_int8_range: Tuple of integers representing `(min, max)`
             range values for all arrays without a specified range. Intended for
             experimenting with quantization via "dummy quantization". (default None)
@@ -97,35 +216,37 @@ def convert_keras_model(
         raise ValueError(
             f"Expected `model` argument to be a `tf.keras.Model` instance, got `{model}`."
         )
-    if inference_input_type not in (tf.float32, tf.int8):
-        raise ValueError(
-            "Expected `inference_input_type` to be either `tf.float32` or `tf.int8`, "
-            f"got {inference_input_type}."
-        )
-    if inference_output_type not in (tf.float32, tf.int8):
-        raise ValueError(
-            "Expected `inference_output_type` to be either `tf.float32` or `tf.int8`, "
-            f"got {inference_output_type}."
-        )
-    if target not in ("arm", "xcore"):
-        raise ValueError(f'Expected `target` to be "arm" or "xcore", but got {target}.')
-
-    if not tf.executing_eagerly():
-        raise RuntimeError(
-            "Graph mode is not supported. Please enable eager execution using "
-            "tf.enable_eager_execution() when using TensorFlow 1.x"
-        )
-    if experimental_default_int8_range:
-        warnings.warn(
-            "Using `experimental_default_int8_range` as fallback quantization stats. "
-            "This should only be used for latency tests."
-        )
     if hasattr(model, "dtype_policy") and model.dtype_policy.name != "float32":
-        raise RuntimeError(
+        raise ValueError(
             "Mixed precision float16 models are not supported by the TFLite converter, "
             "please convert them to float32 first. See also: "
             "https://github.com/tensorflow/tensorflow/issues/46380"
         )
+    _validate_options(
+        inference_input_type=inference_input_type,
+        inference_output_type=inference_output_type,
+        target=target,
+        experimental_default_int8_range=experimental_default_int8_range,
+    )
+
+    # First attempt conversion as saved model
+    try:
+        with tempfile.TemporaryDirectory() as saved_model_dir:
+            model.save(saved_model_dir, save_format="tf")
+
+            return convert_saved_model(
+                saved_model_dir,
+                inference_input_type=inference_input_type,
+                inference_output_type=inference_output_type,
+                experimental_default_int8_range=experimental_default_int8_range,
+                experimental_enable_bitpacked_activations=experimental_enable_bitpacked_activations,
+                target=target,
+            )
+    except Exception:
+        warnings.warn(
+            "Saved-model conversion failed, falling back to graphdef-based conversion."
+        )
+
     func = concrete_function_from_keras_model(model)
     if version.parse(tf.__version__) >= version.parse("1.15"):
         frozen_func = convert_variables_to_constants_v2(func, lower_control_flow=False)
@@ -168,6 +289,7 @@ def convert_keras_model(
         experimental_default_int8_range,
         experimental_enable_bitpacked_activations,
     )
+
     if should_quantize and (
         inference_input_type != tf.float32 or inference_output_type != tf.float32
     ):
diff --git a/larq_compute_engine/mlir/python/converter_test.py b/larq_compute_engine/mlir/python/converter_test.py
@@ -1,21 +1,24 @@
 import sys
 import unittest
+from packaging import version
 from unittest import mock
 
+import tensorflow as tf
 import larq_zoo as lqz
 from tensorflow.python.eager import context
 
 sys.modules["importlib.metadata"] = mock.MagicMock()
 sys.modules["importlib_metadata"] = mock.MagicMock()
-sys.modules["larq_compute_engine.mlir._graphdef_tfl_flatbuffer"] = mock.MagicMock()
+sys.modules["larq_compute_engine.mlir._tf_tfl_flatbuffer"] = mock.MagicMock()
 sys.modules[
     "larq_compute_engine.tflite.python.interpreter_wrapper_lite"
 ] = mock.MagicMock()
 sys.modules["larq_compute_engine.mlir.python.tflite_schema"] = mock.MagicMock()
 
 from larq_compute_engine.mlir.python.converter import convert_keras_model
-from larq_compute_engine.mlir._graphdef_tfl_flatbuffer import (
-    convert_graphdef_to_tflite_flatbuffer as mocked_converter,
+from larq_compute_engine.mlir._tf_tfl_flatbuffer import (
+    convert_graphdef_to_tflite_flatbuffer as mocked_graphdef_converter,
+    convert_saved_model_to_tflite_flatbuffer as mocked_saved_model_converter,
 )
 
 
@@ -24,17 +27,22 @@ def test_larq_zoo_models(self):
         with context.eager_mode():
             model = lqz.sota.QuickNet(weights=None)
             convert_keras_model(model)
-        mocked_converter.assert_called_once_with(
-            mock.ANY,
-            ["input_1"],
-            ["DT_FLOAT"],
-            [[1, 224, 224, 3]],
-            ["Identity"],
-            False,
-            "arm",
-            None,
-            False,
-        )
+        if version.parse(tf.__version__) < version.parse("2.2"):
+            mocked_graphdef_converter.assert_called_once_with(
+                mock.ANY,
+                ["input_1"],
+                ["DT_FLOAT"],
+                [[1, 224, 224, 3]],
+                ["Identity"],
+                False,
+                "arm",
+                None,
+                False,
+            )
+        else:
+            mocked_saved_model_converter.assert_called_once_with(
+                mock.ANY, ["serve"], ["serving_default"], 1, "arm", None, False
+            )
 
     def test_wrong_arg(self):
         with self.assertRaises(ValueError):
diff --git a/larq_compute_engine/mlir/python/graphdef_tfl_flatbuffer.cc b/larq_compute_engine/mlir/python/graphdef_tfl_flatbuffer.cc
@@ -12,7 +12,6 @@
 #include "pybind11/stl.h"
 #include "tensorflow/compiler/mlir/lite/quantization/quantization_config.h"
 #include "tensorflow/compiler/mlir/lite/transforms/passes.h"
-#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
 #include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
 #include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
 #include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
@@ -118,8 +117,3 @@ pybind11::bytes ConvertGraphDefToTFLiteFlatBuffer(
 }
 
 }  // namespace tensorflow
-
-PYBIND11_MODULE(_graphdef_tfl_flatbuffer, m) {
-  m.def("convert_graphdef_to_tflite_flatbuffer",
-        &tensorflow::ConvertGraphDefToTFLiteFlatBuffer);
-};
diff --git a/larq_compute_engine/mlir/python/pybind_export.cc b/larq_compute_engine/mlir/python/pybind_export.cc
@@ -0,0 +1,32 @@
+#include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
+#include "pybind11/stl.h"
+
+namespace tensorflow {
+
+using std::string;
+
+pybind11::bytes ConvertGraphDefToTFLiteFlatBuffer(
+    const pybind11::bytes& graphdef_bytes,
+    const std::vector<string>& input_arrays,
+    const std::vector<string>& input_dtypes,
+    const std::vector<std::vector<int>>& input_shapes,
+    const std::vector<string>& output_arrays, const bool should_quantize,
+    const std::string& target_str, const pybind11::object& default_ranges,
+    const bool experimental_enable_bitpacked_activations);
+
+pybind11::bytes ConvertSavedModelToTFLiteFlatBuffer(
+    const std::string& saved_model_dir,
+    const std::vector<std::string>& saved_model_tags,
+    const std::vector<std::string>& exported_names,
+    const int saved_model_version, const std::string& target_str,
+    const pybind11::object& default_ranges,
+    const bool experimental_enable_bitpacked_activations);
+}  // namespace tensorflow
+
+PYBIND11_MODULE(_tf_tfl_flatbuffer, m) {
+  m.def("convert_graphdef_to_tflite_flatbuffer",
+        &tensorflow::ConvertGraphDefToTFLiteFlatBuffer);
+  m.def("convert_saved_model_to_tflite_flatbuffer",
+        &tensorflow::ConvertSavedModelToTFLiteFlatBuffer);
+};
diff --git a/larq_compute_engine/mlir/python/saved_model_tfl_flatbuffer.cc b/larq_compute_engine/mlir/python/saved_model_tfl_flatbuffer.cc
diff --git a/larq_compute_engine/mlir/tf_tfl_passes.cc b/larq_compute_engine/mlir/tf_tfl_passes.cc
diff --git a/larq_compute_engine/tests/end2end_test.py b/larq_compute_engine/tests/end2end_test.py
