Add N-bit QAT features (fake-quantization).

PiperOrigin-RevId: 395863467

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/api/BUILD

@@ -17,6 +17,9 @@ py_strict_library(
         "quantization/keras/collaborative_optimizations/__init__.py",
         "quantization/keras/default_8bit/__init__.py",
         "quantization/keras/default_8bit/default_8bit_transforms/__init__.py",
+        "quantization/keras/experimental/__init__.py",
+        "quantization/keras/experimental/default_n_bit/__init__.py",
+        "quantization/keras/experimental/default_n_bit/default_n_bit_transforms/__init__.py",
         "quantization/keras/graph_transformations/__init__.py",
         "quantization/keras/graph_transformations/model_transformer/__init__.py",
         "quantization/keras/graph_transformations/transforms/__init__.py",
@@ -46,6 +49,10 @@ py_strict_library(
         "//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/default_8bit:default_8bit_transforms",
+        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit:default_n_bit_quantize_layout_transform",
+        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit:default_n_bit_quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit:default_n_bit_quantize_scheme",
+        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit:default_n_bit_transforms",
         "//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:model_transformer",
         "//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:transforms",
         "//tensorflow_model_optimization/python/core/sparsity/keras:prunable_layer",

tensorflow_model_optimization/python/core/api/quantization/keras/__init__.py

@@ -20,6 +20,7 @@
 from tensorflow_model_optimization.python.core.api.quantization.keras import default_8bit
 from tensorflow_model_optimization.python.core.api.quantization.keras import graph_transformations
 from tensorflow_model_optimization.python.core.api.quantization.keras import collaborative_optimizations
+from tensorflow_model_optimization.python.core.api.quantization.keras import experimental
 
 # quantize all layers with default quantization implementation.
 from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_model

tensorflow_model_optimization/python/core/api/quantization/keras/experimental/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+# ==============================================================================
+"""Module containing experimental quantization features."""
+# pylint: disable=g-bad-import-order
+
+# submodules
+from tensorflow_model_optimization.python.core.api.quantization.keras.experimental import default_n_bit
+
+# pylint: enable=g-bad-import-order

tensorflow_model_optimization/python/core/api/quantization/keras/experimental/default_n_bit/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+# ==============================================================================
+"""Module containing N-bit default quantization scheme."""
+# pylint: disable=g-bad-import-order
+
+# submodules
+from tensorflow_model_optimization.python.core.api.quantization.keras.experimental.default_n_bit import default_n_bit_transforms
+
+# The N-bit default quantization scheme classes.
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_quantize_scheme import DefaultNBitQuantizeScheme
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_quantize_layout_transform import DefaultNBitQuantizeLayoutTransform
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_quantize_registry import DefaultNBitQuantizeRegistry
+
+# pylint: enable=g-bad-import-order

tensorflow_model_optimization/python/core/api/quantization/keras/experimental/default_n_bit/default_n_bit_transforms/__init__.py

@@ -0,0 +1,33 @@
+# 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.
+# ==============================================================================
+"""Module containing N-bit default transforms."""
+
+# The 8bit default transform classes.
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import ConcatTransform
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import ConcatTransform3Inputs
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import ConcatTransform4Inputs
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import ConcatTransform5Inputs
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import ConcatTransform6Inputs
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DBatchNormActivationQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DBatchNormQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DBatchNormReLUQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DReshapeBatchNormActivationQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DReshapeBatchNormQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DReshapeBatchNormReLUQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import InputLayerQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import LayerReluActivationQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import LayerReLUQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import SeparableConv1DQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import SeparableConvQuantize

tensorflow_model_optimization/python/core/quantization/keras/BUILD

@@ -19,6 +19,7 @@ py_strict_library(
         "//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations",  # buildcleaner: keep
         "//tensorflow_model_optimization/python/core/quantization/keras/layers",  # buildcleaner: keep
         "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit",  # buildcleaner: keep
+        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit",  # buildcleaner: keep
         "//tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations",  # buildcleaner: keep
     ],
 )
@@ -28,6 +29,7 @@ py_strict_library(
     srcs = ["quant_ops.py"],
     srcs_version = "PY3",
     deps = [
+        # absl/logging dep1,
         # tensorflow dep1,
         # python:training tensorflow dep2,
         "//tensorflow_model_optimization/python/core/keras:compat",
@@ -252,6 +254,7 @@ py_strict_library(
         "//tensorflow_model_optimization/python/core/keras:metrics",
         "//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/experimental/default_n_bit:default_n_bit_quantize_registry",
     ],
 )
 

tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/BUILD

@@ -0,0 +1,152 @@
+load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library", "py_strict_test")
+
+package(default_visibility = [
+    "//tensorflow_model_optimization:__subpackages__",
+])
+
+licenses(["notice"])
+
+py_strict_library(
+    name = "default_n_bit",
+    srcs = [
+        "__init__.py",
+    ],
+    srcs_version = "PY3",
+    deps = [],
+)
+
+py_strict_library(
+    name = "default_n_bit_quantizers",
+    srcs = [
+        "default_n_bit_quantizers.py",
+    ],
+    srcs_version = "PY3",
+    deps = [
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
+    ],
+)
+
+py_test(
+    name = "default_n_bit_quantizers_test",
+    srcs = [
+        "default_n_bit_quantizers_test.py",
+    ],
+    python_version = "PY3",
+    deps = [
+        ":default_n_bit_quantizers",
+        # absl/testing:parameterized dep1,
+        # tensorflow dep1,
+    ],
+)
+
+py_strict_library(
+    name = "default_n_bit_quantize_configs",
+    srcs = [
+        "default_n_bit_quantize_configs.py",
+    ],
+    srcs_version = "PY3",
+    deps = [
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_config",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
+    ],
+)
+
+py_strict_library(
+    name = "default_n_bit_quantize_registry",
+    srcs = [
+        "default_n_bit_quantize_registry.py",
+    ],
+    srcs_version = "PY3",
+    deps = [
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_config",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
+        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit:default_n_bit_quantize_configs",
+        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit:default_n_bit_quantizers",
+    ],
+)
+
+py_test(
+    name = "default_n_bit_quantize_registry_test",
+    srcs = [
+        "default_n_bit_quantize_registry_test.py",
+    ],
+    python_version = "PY3",
+    deps = [
+        ":default_n_bit_quantize_registry",
+        # absl/testing:parameterized dep1,
+        # numpy dep1,
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
+    ],
+)
+
+py_library(
+    name = "default_n_bit_transforms",
+    srcs = [
+        "default_n_bit_transforms.py",
+    ],
+    srcs_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        # numpy dep1,
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_aware_activation",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_layer",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
+        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit:default_n_bit_quantize_configs",
+        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit:default_n_bit_quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:transforms",
+    ],
+)
+
+py_strict_library(
+    name = "default_n_bit_quantize_layout_transform",
+    srcs = [
+        "default_n_bit_quantize_layout_transform.py",
+    ],
+    srcs_version = "PY3",
+    deps = [
+        ":default_n_bit_transforms",
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_layout_transform",
+        "//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:model_transformer",
+    ],
+)
+
+py_strict_test(
+    name = "default_n_bit_transforms_test",
+    size = "large",
+    srcs = [
+        "default_n_bit_transforms_test.py",
+    ],
+    python_version = "PY3",
+    deps = [
+        ":default_n_bit_quantize_configs",
+        ":default_n_bit_transforms",
+        # absl/testing:parameterized dep1,
+        # numpy dep1,
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_aware_activation",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_layer",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
+        "//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:model_transformer",
+        "//tensorflow_model_optimization/python/core/quantization/keras/layers:conv_batchnorm_test_utils",
+    ],
+)
+
+py_strict_library(
+    name = "default_n_bit_quantize_scheme",
+    srcs = [
+        "default_n_bit_quantize_scheme.py",
+    ],
+    srcs_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":default_n_bit_quantize_layout_transform",
+        ":default_n_bit_quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_scheme",
+    ],
+)

tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/README.md

@@ -0,0 +1,2 @@
+This directory is modified based on default_8bit, which allows you to manually
+change the number of bits of weight and activation in QAT.

tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2019 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.
+# ==============================================================================

tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantize_configs.py

@@ -0,0 +1,72 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Default N-Bit QuantizeConfigs."""
+
+from typing import Any, Dict
+from tensorflow_model_optimization.python.core.quantization.keras import quantize_config
+from tensorflow_model_optimization.python.core.quantization.keras import quantizers
+
+
+class DefaultNBitOutputQuantizeConfig(quantize_config.QuantizeConfig):
+  """QuantizeConfig which only quantizes the output from a layer."""
+
+  def __init__(self, num_bits_weight: int = 8, num_bits_activation: int = 8):
+    self._num_bits_weight = num_bits_weight
+    self._num_bits_activation = num_bits_activation
+
+  def get_weights_and_quantizers(self, layer):
+    return []
+
+  def get_activations_and_quantizers(self, layer):
+    return []
+
+  def set_quantize_weights(self, layer, quantize_weights):
+    pass
+
+  def set_quantize_activations(self, layer, quantize_activations):
+    pass
+
+  def get_output_quantizers(self, layer):
+    return [quantizers.MovingAverageQuantizer(
+        num_bits=self._num_bits_activation, per_axis=False,
+        symmetric=False, narrow_range=False)]  # activation/output
+
+  def get_config(self) -> Dict[str, Any]:
+    return {
+        'num_bits_weight': self._num_bits_weight,
+        'num_bits_activation': self._num_bits_activation,
+    }
+
+
+class NoOpQuantizeConfig(quantize_config.QuantizeConfig):
+  """QuantizeConfig which does not quantize any part of the layer."""
+
+  def get_weights_and_quantizers(self, layer):
+    return []
+
+  def get_activations_and_quantizers(self, layer):
+    return []
+
+  def set_quantize_weights(self, layer, quantize_weights):
+    pass
+
+  def set_quantize_activations(self, layer, quantize_activations):
+    pass
+
+  def get_output_quantizers(self, layer):
+    return []
+
+  def get_config(self):
+    return {}