Add weight clustering implementation using compression API.

It requires clustering modules in tfmot weight clustering implementaion.

Added dtype for init variable due to a compressed weight dtype is int64.
On example code, gzip compressed tflite file size is reduced from 1592595 bytes to 224167 bytes. (14%)

PiperOrigin-RevId: 341767464

Author: Xhark

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/BUILD

@@ -67,3 +67,26 @@ py_test(
         # tensorflow dep1,
     ],
 )
+
+py_library(
+    name = "weight_clustering",
+    srcs = ["weight_clustering.py"],
+    srcs_version = "PY3",
+    deps = [
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/clustering/keras:clustering_centroids",
+        "//tensorflow_model_optimization/python/core/clustering/keras:clustering_registry",
+        "//tensorflow_model_optimization/python/core/common/keras/compression:algorithm",
+    ],
+)
+
+py_test(
+    name = "weight_clustering_test",
+    srcs = ["weight_clustering_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":weight_clustering",
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config",
+    ],
+)

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/weight_clustering.py

@@ -0,0 +1,120 @@
+# 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.
+# ==============================================================================
+"""Weight clustering algorithm using tfmot compression api."""
+from typing import List
+
+import tensorflow as tf
+
+# TODO(tfmot): Make sure weight clustering APIs can be used in this place or
+# move the APIs into the same directory.
+from tensorflow_model_optimization.python.core.clustering.keras import clustering_centroids
+from tensorflow_model_optimization.python.core.clustering.keras import clustering_registry
+from tensorflow_model_optimization.python.core.common.keras.compression import algorithm
+
+
+class WeightClusteringParams(object):
+  """Weight clustering parameters."""
+
+  def __init__(self,
+               number_of_clusters,
+               cluster_centroids_init):
+    self.number_of_clusters = number_of_clusters
+    self.cluster_centroids_init = cluster_centroids_init
+
+
+class WeightClustering(algorithm.WeightCompressionAlgorithm):
+  """Weight clustering compression module config."""
+
+  def __init__(self, params):
+    self.params = params
+
+  def init_training_weights_repr(
+      self, pretrained_weight: tf.Tensor) -> List[algorithm.WeightRepr]:
+    """Init function from pre-trained model case."""
+    centroid_initializer = clustering_centroids.CentroidsInitializerFactory.\
+        get_centroid_initializer(
+            self.params.cluster_centroids_init
+        )(pretrained_weight, self.params.number_of_clusters)
+
+    cluster_centroids = centroid_initializer.get_cluster_centroids()
+
+    if len(pretrained_weight.shape) == 2:
+      clustering_impl_cls = clustering_registry.DenseWeightsCA
+    elif len(pretrained_weight.shape) == 4:
+      clustering_impl_cls = clustering_registry.ConvolutionalWeightsCA
+    else:
+      raise NotImplementedError('Only for dimension=2 or 4 is supported.')
+
+    clustering_impl = clustering_impl_cls(
+        cluster_centroids
+    )
+
+    # We find the nearest cluster centroids and store them so that ops can
+    # build their weights upon it. These indices are calculated once and
+    # stored forever. We use to make look-ups from self.cluster_centroids_tf
+    pulling_indices = clustering_impl.get_pulling_indices(pretrained_weight)
+
+    return [
+        algorithm.WeightRepr(
+            name='cluster_centroids',
+            shape=cluster_centroids.shape,
+            dtype=cluster_centroids.dtype,
+            initializer=tf.keras.initializers.Constant(cluster_centroids)),
+        algorithm.WeightRepr(
+            name='pulling_indices',
+            shape=pulling_indices.shape,
+            dtype=pulling_indices.dtype,
+            initializer=tf.keras.initializers.Constant(pulling_indices))
+    ]
+
+  def decompress(self,
+                 cluster_centroids: tf.Tensor,
+                 pulling_indices: tf.Tensor) -> tf.Tensor:
+    return tf.reshape(
+        tf.gather(cluster_centroids,
+                  tf.reshape(pulling_indices, shape=(-1,))),
+        pulling_indices.shape)
+
+  def training(self,
+               cluster_centroids: tf.Tensor,
+               pulling_indices: tf.Tensor) -> tf.Tensor:
+    return self.decompress(cluster_centroids, pulling_indices)
+
+  def get_compressible_weights(
+      self, original_layer: tf.keras.layers.Layer) -> List[str]:
+    if isinstance(original_layer, tf.keras.layers.Conv2D) or \
+       isinstance(original_layer, tf.keras.layers.Dense):
+      return ['kernel']
+    return []
+
+
+def optimize(
+    to_optimize: tf.keras.Model,
+    params: WeightClusteringParams) -> tf.keras.Model:
+  """Model developer API for optimizing a model."""
+
+  def _optimize_layer(layer):
+    # Require layer to be built so that the SVD-factorized weights
+    # can be initialized from the weights.
+    if not layer.built:
+      raise ValueError(
+          'Applying weight clustering currently '
+          'requires passing in a built model')
+
+    return algorithm.create_layer_for_training(
+        layer, algorithm=WeightClustering(params))
+
+  return tf.keras.models.clone_model(
+      to_optimize, clone_function=_optimize_layer)

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/weight_clustering_test.py

@@ -0,0 +1,158 @@
+# 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.
+# ==============================================================================
+"""Tests for weight clustering algorithm."""
+
+import os
+import tempfile
+
+import tensorflow as tf
+
+from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
+from tensorflow_model_optimization.python.core.common.keras.compression.algorithms import weight_clustering
+
+
+def _build_model():
+  i = tf.keras.layers.Input(shape=(28, 28), name='input')
+  x = tf.keras.layers.Reshape((28, 28, 1))(i)
+  x = tf.keras.layers.Conv2D(
+      20, 5, activation='relu', padding='valid', name='conv1')(
+          x)
+  x = tf.keras.layers.MaxPool2D(2, 2)(x)
+  x = tf.keras.layers.Conv2D(
+      50, 5, activation='relu', padding='valid', name='conv2')(
+          x)
+  x = tf.keras.layers.MaxPool2D(2, 2)(x)
+  x = tf.keras.layers.Flatten()(x)
+  x = tf.keras.layers.Dense(500, activation='relu', name='fc1')(x)
+  output = tf.keras.layers.Dense(10, name='fc2')(x)
+
+  model = tf.keras.Model(inputs=[i], outputs=[output])
+  return model
+
+
+def _get_dataset():
+  mnist = tf.keras.datasets.mnist
+  (x_train, y_train), (x_test, y_test) = mnist.load_data()
+  x_train, x_test = x_train / 255.0, x_test / 255.0
+  # Use subset of 60000 examples to keep unit test speed fast.
+  x_train = x_train[:1000]
+  y_train = y_train[:1000]
+
+  return (x_train, y_train), (x_test, y_test)
+
+
+def _train_model(model):
+  loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+  model.compile(optimizer='adam', loss=loss_fn, metrics=['accuracy'])
+  (x_train, y_train), _ = _get_dataset()
+  model.fit(x_train, y_train, epochs=1)
+
+
+def _save_as_saved_model(model):
+  saved_model_dir = tempfile.mkdtemp()
+  model.save(saved_model_dir)
+  return saved_model_dir
+
+
+def _get_directory_size_in_bytes(directory):
+  total = 0
+  try:
+    for entry in os.scandir(directory):
+      if entry.is_file():
+        # if it's a file, use stat() function
+        total += entry.stat().st_size
+      elif entry.is_dir():
+        # if it's a directory, recursively call this function
+        total += _get_directory_size_in_bytes(entry.path)
+  except NotADirectoryError:
+    # if `directory` isn't a directory, get the file size then
+    return os.path.getsize(directory)
+  except PermissionError:
+    # if for whatever reason we can't open the folder, return 0
+    return 0
+  return total
+
+
+class FunctionalTest(tf.test.TestCase):
+
+  def testWeightClustering_TrainingE2E(self):
+    number_of_clusters = 8
+    model = _build_model()
+    _train_model(model)
+    original_saved_model_dir = _save_as_saved_model(model)
+
+    params = weight_clustering.WeightClusteringParams(
+        number_of_clusters=number_of_clusters,
+        cluster_centroids_init=\
+        cluster_config.CentroidInitialization.DENSITY_BASED)
+    compressed_model = weight_clustering.optimize(model, params)
+
+    _train_model(compressed_model)
+
+    saved_model_dir = _save_as_saved_model(compressed_model)
+
+    _, (x_test, y_test) = _get_dataset()
+
+    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+
+    compressed_model.compile(
+        optimizer='adam', loss=loss_fn, metrics=['accuracy'])
+
+    results = compressed_model.evaluate(x_test, y_test)
+
+    # Accuracy test.
+    self.assertGreater(results[1], 0.85)  # 0.8708
+
+    original_size = _get_directory_size_in_bytes(original_saved_model_dir)
+    compressed_size = _get_directory_size_in_bytes(saved_model_dir)
+
+    # Compressed model size test.
+    # TODO(tfmot): gzip compression can reduce file size much better.
+    self.assertLess(compressed_size, original_size / 1.3)
+
+  def testWeightClustering_SingleLayer(self):
+    number_of_clusters = 8
+    i = tf.keras.layers.Input(shape=(2), name='input')
+    output = tf.keras.layers.Dense(3, name='fc1')(i)
+    model = tf.keras.Model(inputs=[i], outputs=[output])
+
+    dense_layer_weights = model.layers[1].get_weights()
+
+    params = weight_clustering.WeightClusteringParams(
+        number_of_clusters=number_of_clusters,
+        cluster_centroids_init=\
+        cluster_config.CentroidInitialization.DENSITY_BASED)
+    compressed_model = weight_clustering.optimize(model, params)
+
+    dense_layer_compressed_weights = compressed_model.layers[1].get_weights()
+
+    # clustering_centroids.
+    self.assertEqual(
+        dense_layer_compressed_weights[0].shape, (number_of_clusters,))
+
+    # pulling_indices.
+    self.assertEqual(
+        dense_layer_compressed_weights[1].shape,
+        dense_layer_weights[0].shape)
+    self.assertEqual(str(dense_layer_compressed_weights[1].dtype), 'int64')
+    self.assertAllInRange(
+        dense_layer_compressed_weights[1], 0, number_of_clusters - 1)
+
+    # bias
+    assert (dense_layer_weights[1] == dense_layer_compressed_weights[2]).all()
+
+
+if __name__ == '__main__':
+  tf.test.main()