Expose QuantizationScheme in the QAT API.

PiperOrigin-RevId: 341923968

Author: Xhark

tensorflow_model_optimization/python/core/quantization/keras/BUILD

@@ -243,8 +243,8 @@ py_library(
         ":quantize_layer",
         ":quantize_wrapper",
         # tensorflow dep1,
-        "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_layout_transform",
         "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_scheme",
         "//tensorflow_model_optimization/python/core/quantization/keras/layers:conv_batchnorm",
     ],
 )
@@ -329,3 +329,12 @@ py_library(
         # tensorflow dep1,
     ],
 )
+
+py_library(
+    name = "quantize_scheme",
+    srcs = [
+        "quantize_scheme.py",
+    ],
+    srcs_version = "PY3",
+    visibility = ["//visibility:public"],
+)

tensorflow_model_optimization/python/core/quantization/keras/default_8bit/BUILD

@@ -147,3 +147,16 @@ py_test(
         "//tensorflow_model_optimization/python/core/quantization/keras:utils",
     ],
 )
+
+py_library(
+    name = "default_8bit_quantize_scheme",
+    srcs = [
+        "default_8bit_quantize_scheme.py",
+    ],
+    srcs_version = "PY3",
+    deps = [
+        ":default_8bit_quantize_layout_transform",
+        ":default_8bit_quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_scheme",
+    ],
+)

tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_quantize_scheme.py

@@ -0,0 +1,29 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Quantization scheme which specifies how quantization should be applied."""
+
+from tensorflow_model_optimization.python.core.quantization.keras import quantize_scheme
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_layout_transform
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_registry
+
+
+class Default8BitQuantizeScheme(quantize_scheme.QuantizeScheme):
+
+  def get_layout_transformer(self):
+    return default_8bit_quantize_layout_transform.QuantizeLayoutTransform()
+
+  def get_quantize_registry(self):
+    return default_8bit_quantize_registry.QuantizeRegistry()
+

tensorflow_model_optimization/python/core/quantization/keras/quantize.py

@@ -22,8 +22,8 @@
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_layer
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_wrapper
 from tensorflow_model_optimization.python.core.quantization.keras import quantizers
-from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_layout_transform
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_registry
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_scheme
 from tensorflow_model_optimization.python.core.quantization.keras.layers import conv_batchnorm
 
 keras = tf.keras
@@ -263,7 +263,9 @@ def quantize_annotate_layer(to_annotate, quantize_config=None):
       layer=to_annotate, quantize_config=quantize_config)
 
 
-def quantize_apply(model):
+def quantize_apply(
+    model,
+    scheme=default_8bit_quantize_scheme.Default8BitQuantizeScheme()):
   """Quantize a `tf.keras` model that has been annotated for quantization.
 
   Quantization constructs a model which emulates quantization during training.
@@ -298,6 +300,8 @@ def quantize_apply(model):
   Args:
     model: A `tf.keras` Sequential or Functional model which has been annotated
       with `quantize_annotate`. It can have pre-trained weights.
+    scheme: A `QuantizeScheme` which specifies transformer and quantization
+      registry. The default is `Default8BitQuantizeScheme()`.
 
   Returns:
     Returns a new `tf.keras` model in which the annotated layers have been
@@ -403,15 +407,13 @@ def _quantize(layer):  # pylint: disable=missing-docstring
 
   # 3. Apply the graph transformations required to match model passes on
   # target device/dialect.
-  quantize_transform = \
-    default_8bit_quantize_layout_transform.QuantizeLayoutTransform()
+  quantize_transform = scheme.get_layout_transformer()
   # layer_quantize_map gets modified by the transformations.
   transformed_model, layer_quantize_map = quantize_transform.apply(
       unwrapped_model, layer_quantize_map)
 
   # TODO(pulkitb): Think more about how to introduce Default specific code.
-  quantize_registry = default_8bit_quantize_registry.QuantizeRegistry(
-  )
+  quantize_registry = scheme.get_quantize_registry()
 
   # 4. Actually quantize all the relevant layers in the model. This is done by
   # wrapping the layers with QuantizeWrapper, and passing the associated

tensorflow_model_optimization/python/core/quantization/keras/quantize_scheme.py

@@ -0,0 +1,44 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Quantization scheme which specifies how quantization should be applied.
+
+Module: tfmot.quantization.keras
+"""
+
+import abc
+import six
+
+
+@six.add_metaclass(abc.ABCMeta)
+class QuantizeScheme(object):
+  """ABC interface which specifies transformer and quantization registry."""
+
+  @abc.abstractmethod
+  def get_layout_transformer(self):
+    """Returns the layout transforms for this scheme.
+
+    Returns:
+      Returns the QuantizeLayoutTransform for this quantization scheme.
+    """
+    raise NotImplementedError('Must be implemented in subclasses.')
+
+  @abc.abstractmethod
+  def get_quantize_registry(self):
+    """Returns the quantization registry for this scheme.
+
+    Returns:
+      Returns the QuantizeRegistry for this quantization scheme.
+    """
+    raise NotImplementedError('Must be implemented in subclasses.')