[Collaborative_optimizations] Cluster-preserving quantization aware training

Author: Ruomei

tensorflow_model_optimization/python/core/api/experimental/combine/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 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.
@@ -13,5 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 """Module containing collaborative optimization code."""
+
 from tensorflow_model_optimization.python.core.quantization.keras.collaborative_optimizations.prune_preserve.default_8bit_prune_preserve_quantize_scheme import (
     Default8BitPrunePreserveQuantizeScheme,)
+
+from tensorflow_model_optimization.python.core.quantization.keras.collaborative_optimizations.cluster_preserve.default_8bit_cluster_preserve_quantize_scheme import (
+    Default8BitClusterPreserveQuantizeScheme,)
+
+from tensorflow_model_optimization.python.core.quantization.keras.collaborative_optimizations.cluster_preserve.cluster_utils import (
+    strip_clustering_cqat,)

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/BUILD

@@ -6,5 +6,8 @@ py_library(
     name = "collaborative_optimizations",
     srcs = ["__init__.py"],
     srcs_version = "PY3",
-    deps = ["//tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/prune_preserve"],
+    deps = [
+        "//tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/prune_preserve",
+        "//tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve",
+    ],
 )

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve/BUILD

@@ -0,0 +1,72 @@
+package(default_visibility = [
+    "//tensorflow_model_optimization:__subpackages__",
+])
+
+licenses(["notice"])  # Apache 2.0
+
+py_library(
+    name = "cluster_preserve",
+    srcs = [
+        "__init__.py",
+    ],
+    srcs_version = "PY3",
+    deps = [
+        ":default_8bit_cluster_preserve_quantize_scheme",
+    ],
+)
+
+py_library(
+    name = "cluster_utils",
+    srcs = [
+        "cluster_utils.py",
+    ],
+    srcs_version = "PY3",
+    visibility = ["//visibility:private"],
+    deps = [
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/clustering/keras:clustering_registry",
+    ],
+)
+
+py_library(
+    name = "cluster_preserve_quantize_registry",
+    srcs = [
+        "cluster_preserve_quantize_registry.py",
+    ],
+    srcs_version = "PY3",
+    visibility = ["//visibility:private"],
+    deps = [
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantizers",
+        "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_registry",
+        ":cluster_utils",
+    ],
+)
+
+py_test(
+    name = "cluster_preserve_quantize_registry_test",
+    srcs = [
+        "cluster_preserve_quantize_registry_test.py",
+    ],
+    python_version = "PY3",
+    visibility = ["//visibility:private"],
+    deps = [
+        # tensorflow dep1,
+        ":cluster_preserve_quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_config",
+    ]
+)
+
+py_library(
+    name = "default_8bit_cluster_preserve_quantize_scheme",
+    srcs = [
+        "default_8bit_cluster_preserve_quantize_scheme.py",
+    ],
+    srcs_version = "PY3",
+    visibility = ["//visibility:private"],
+    deps = [
+        ":cluster_preserve_quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_scheme",
+    ],
+)

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve/__init__.py

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