Merge branch 'master' into clusterable_layer

Author: wwwind

tensorflow_model_optimization/python/core/api/BUILD

@@ -13,6 +13,10 @@ py_library(
         "quantization/__init__.py",
         "quantization/keras/__init__.py",
         "quantization/keras/default_8bit/__init__.py",
+        "quantization/keras/default_8bit/default_8bit_transforms/__init__.py",
+        "quantization/keras/graph_transformations/__init__.py",
+        "quantization/keras/graph_transformations/model_transformer/__init__.py",
+        "quantization/keras/graph_transformations/transforms/__init__.py",
         "quantization/keras/quantizers/__init__.py",
         "sparsity/__init__.py",
         "sparsity/keras/__init__.py",

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

@@ -18,6 +18,7 @@
 # submodules
 from tensorflow_model_optimization.python.core.api.quantization.keras import quantizers
 from tensorflow_model_optimization.python.core.api.quantization.keras import default_8bit
+from tensorflow_model_optimization.python.core.api.quantization.keras import graph_transformations
 
 # quantize all layers with default quantization implementation.
 from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_model
@@ -30,6 +31,7 @@
 # quantize with custom quantization parameterization or implementation, or
 # handle custom Keras layers.
 from tensorflow_model_optimization.python.core.quantization.keras.quantize_config import QuantizeConfig
+from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapper import QuantizeWrapper
 
 # Deserialize quantized model for Keras h5 format.
 from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_scope

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

@@ -15,6 +15,9 @@
 """Module containing 8bit default quantization scheme."""
 # pylint: disable=g-bad-import-order
 
+# submodules
+from tensorflow_model_optimization.python.core.api.quantization.keras.default_8bit import default_8bit_transforms
+
 # The 8bit default quantization scheme classes.
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_quantize_scheme import Default8BitQuantizeScheme
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_quantize_layout_transform import Default8BitQuantizeLayoutTransform

tensorflow_model_optimization/python/core/api/quantization/keras/default_8bit/default_8bit_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 8bit default transforms."""
+
+# The 8bit default transform classes.
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import ConcatTransform
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import ConcatTransform3Inputs
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import ConcatTransform4Inputs
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import ConcatTransform5Inputs
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import ConcatTransform6Inputs
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DBatchNormActivationQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DBatchNormQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DBatchNormReLUQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DReshapeBatchNormActivationQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DReshapeBatchNormQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DReshapeBatchNormReLUQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import InputLayerQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import LayerReluActivationQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import LayerReLUQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import SeparableConv1DQuantize
+from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import SeparableConvQuantize

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

@@ -0,0 +1,19 @@
+# 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.
+# ==============================================================================
+"""Module containing code for graph transformation."""
+
+# submodules
+from tensorflow_model_optimization.python.core.api.quantization.keras.graph_transformations import model_transformer
+from tensorflow_model_optimization.python.core.api.quantization.keras.graph_transformations import transforms

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

@@ -0,0 +1,17 @@
+# 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.
+# ==============================================================================
+"""Module containing classes for model transformer."""
+
+from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations.model_transformer import ModelTransformer

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

@@ -0,0 +1,19 @@
+# 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.
+# ==============================================================================
+"""Module containing classes for transform."""
+
+from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations.transforms import LayerNode
+from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations.transforms import LayerPattern
+from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations.transforms import Transform

tensorflow_model_optimization/python/core/clustering/keras/BUILD

@@ -44,7 +44,7 @@ py_library(
     visibility = ["//visibility:public"],
     deps = [
         ":clusterable_layer",
-        ":clustering_algorithm"
+        ":clustering_algorithm",
     ],
 )
 
@@ -183,7 +183,7 @@ py_test(
     python_version = "PY3",
     visibility = ["//visibility:public"],
     deps = [
-        ":cluster"
+        ":cluster",
         # tensorflow dep1,
     ],
 )
@@ -194,7 +194,7 @@ py_test(
     python_version = "PY3",
     visibility = ["//visibility:public"],
     deps = [
-        ":cluster"
+        ":cluster",
         # tensorflow dep1,
     ],
-)
+)

tensorflow_model_optimization/python/core/clustering/keras/cluster_test.py

@@ -54,55 +54,52 @@ def get_clusterable_weights(self):
 class CustomNonClusterableLayer(layers.Dense):
   pass
 
+
 class KerasCustomLayer(keras.layers.Layer):
+
   def __init__(self, units=32):
     super(KerasCustomLayer, self).__init__()
     self.units = units
 
   def build(self, input_shape):
     self.w = self.add_weight(
-      shape=(input_shape[-1], self.units),
-      initializer="random_normal",
-      trainable=True,
+        shape=(input_shape[-1], self.units),
+        initializer='random_normal',
+        trainable=True,
     )
     self.b = self.add_weight(
-      shape=(self.units,),
-      initializer="random_normal",
-      trainable=False
-    )
+        shape=(self.units,), initializer='random_normal', trainable=False)
 
   def call(self, inputs):
     return tf.matmul(inputs, self.w) + self.b
 
-class MyClusterableLayer(keras.layers.Dense,
-  clusterable_layer.ClusterableLayer):
 
-  def __init__(self, num_units):
-      super().__init__(num_units)
+class MyClusterableLayer(keras.layers.Dense,
+                         clusterable_layer.ClusterableLayer):
 
   def get_clusterable_weights(self):
     # Cluster kernel and bias.
     return [('kernel', self.kernel), ('bias', self.bias)]
 
+
 class MyClusterableLayerInvalid(keras.layers.Dense,
-  clusterable_layer.ClusterableLayer):
-  """ This layer is invalid, because it does not provide
-  get_clusterable_weights function.
+                                clusterable_layer.ClusterableLayer):
+  """This layer is invalid: it does not implement get_clusterable_weights(self).
   """
-  def __init__(self, num_units):
-      super().__init__(num_units)
+  pass
+
+
+class TestCustomerableWeightsCA(clustering_registry.AbstractClusteringAlgorithm
+                               ):
+  """Dummy class derived from AbstractClusteringAlgorithm."""
 
-class TestCustomerableWeightsCA(clustering_registry.AbstractClusteringAlgorithm):
-  """ Dummy class derived from AbstractClusteringAlgorithm."""
   def get_pulling_indices(self, weight):
     return [1, 2, 3]
 
+
 class KerasCustomLayerClusterable(keras.layers.Layer,
-  clusterable_layer.ClusterableLayer):
-  """ This keras custom layer is derived from ClusterableLayer
-  and it provides own implementation of the clustering
-  algorithm.
-  """
+                                  clusterable_layer.ClusterableLayer):
+  """Custom Keras clusterable layer, providing its own clustering algorithm."""
 
   def __init__(self):
     super().__init__()
@@ -114,6 +111,7 @@ def get_clusterable_weights(self):
   def get_clusterable_algorithm(self, weight_name):
     return TestCustomerableWeightsCA
 
+
 class ClusterTest(test.TestCase, parameterized.TestCase):
   """Unit tests for the cluster module."""
 
@@ -128,6 +126,7 @@ def setUp(self):
     self.keras_depthwiseconv2d_layer = layers.DepthwiseConv2D((3, 3), (1, 1))
     self.clusterable_layer = MyClusterableLayer(10)
     self.keras_custom_layer = KerasCustomLayer()
+    self.clusterable_layer = MyClusterableLayer(10)
 
 
     clustering_registry.ClusteringLookupRegistry.register_new_implementation(
@@ -274,20 +273,34 @@ def testClusterMyClusterableLayerInvalid(self):
       MyClusterableLayerInvalid(10) # pylint: disable=abstract-class-instantiated
 
   def testClusterKerasCustomLayer(self):
-    """
-    Verifies that attempting to cluster a keras custom layer raises
-    an exception.
-    """
+    """Verifies that attempting to cluster a keras custom layer raises an exception."""
     # If layer is not built, it has not weights, so
     # we just skip it.
     keras_custom_layer = self.keras_custom_layer
-    cluster_wrapper.ClusterWeights(keras_custom_layer,
-                                  **self.params)
+    cluster_wrapper.ClusterWeights(keras_custom_layer, **self.params)
     # We need to build weights before check that clustering is not supported.
     keras_custom_layer.build(input_shape=(10, 10))
     with self.assertRaises(ValueError):
-      cluster_wrapper.ClusterWeights(keras_custom_layer,
-                                     **self.params)
+      cluster_wrapper.ClusterWeights(keras_custom_layer, **self.params)
+
+  def testClusterMyClusterableLayer(self):
+    # we have weights to cluster.
+    layer = self.clusterable_layer
+    layer.build(input_shape=(10, 10))
+
+    wrapped_layer = cluster_wrapper.ClusterWeights(layer, **self.params)
+    self.assertIsInstance(wrapped_layer, cluster_wrapper.ClusterWeights)
+
+  def testKerasCustomLayerClusterable(self):
+    """Verifies that we can wrap keras custom layer that is customerable."""
+    layer = KerasCustomLayerClusterable()
+    wrapped_layer = cluster_wrapper.ClusterWeights(layer, **self.params)
+    self.assertIsInstance(wrapped_layer, cluster_wrapper.ClusterWeights)
+
+  def testClusterMyClusterableLayerInvalid(self):
+    """Verifies that an exception is raised when get_clusterable_weights() is not implemented."""
+    with self.assertRaises(TypeError):
+      MyClusterableLayerInvalid(10)  # pylint: disable=abstract-class-instantiated
 
   @keras_parameterized.run_all_keras_modes
   def testClusterSequentialModelSelectively(self):

tensorflow_model_optimization/python/core/clustering/keras/cluster_wrapper_test.py

@@ -15,6 +15,8 @@
 """Tests for keras ClusterWeights wrapper API."""
 
 import itertools
+import os
+import tempfile
 
 from absl.testing import parameterized
 import tensorflow as tf
@@ -230,63 +232,46 @@ def assert_all_weights_associated(weights, centroid_index):
     # Weights should now be all clustered with the centroid 1
     assert_all_weights_associated(l.layer.kernel, centroid_index=1)
 
-  def testClusterReassociation2(self):
-    """Verifies that the association of weights to cluster centroids are updated every iteration."""
-
+  def testSameWeightsAreReturnedBeforeAndAfterSerialisation(self):
+    """Verify weights of cluster_wrapper are the same after serialisation."""
     # Create a dummy layer for this test
     input_shape = (1, 2,)
-    l = cluster_wrapper.ClusterWeights(
+    original_layer = cluster_wrapper.ClusterWeights(
         keras.layers.Dense(8, input_shape=input_shape),
         number_of_clusters=2,
         cluster_centroids_init=CentroidInitialization.LINEAR
     )
     # Build a layer with the given shape
-    l.build(input_shape)
-
-    # Get name of the clusterable weights
-    clusterable_weights = l.layer.get_clusterable_weights()
-    self.assertLen(clusterable_weights, 1)
-    weights_name = clusterable_weights[0][0]
-    self.assertEqual(weights_name, 'kernel')
-    # Get cluster centroids
-    centroids = l.cluster_centroids_tf[weights_name]
-
-    # Calculate some statistics of the weights to set the centroids later on
-    mean_weight = tf.reduce_mean(l.layer.kernel)
-    min_weight = tf.reduce_min(l.layer.kernel)
-    max_weight = tf.reduce_max(l.layer.kernel)
-    max_dist = max_weight - min_weight
-
-    def assert_all_weights_associated(weights, centroid_index):
-      """Helper function to make sure that all weights are associated with one centroid."""
-      all_associated = tf.reduce_all(
-          tf.equal(
-              weights,
-              tf.constant(centroids[centroid_index], shape=weights.shape)
-          )
-      )
-      self.assertTrue(all_associated)
-
-    # Set centroids so that all weights should be re-associated with centroid 0
-    centroids[0].assign(mean_weight)
-    centroids[1].assign(mean_weight + 2.0 * max_dist)
-
-    # Update associations of weights to centroids
-    l.call(tf.ones(shape=input_shape))
-
-    # Weights should now be all clustered with the centroid 0
-    assert_all_weights_associated(l.layer.kernel, centroid_index=0)
-
-    # Set centroids so that all weights should be re-associated with centroid 1
-    centroids[0].assign(mean_weight - 2.0 * max_dist)
-    centroids[1].assign(mean_weight)
-
-    # Update associations of weights to centroids
-    l.call(tf.ones(shape=input_shape))
-
-    # Weights should now be all clustered with the centroid 1
-    assert_all_weights_associated(l.layer.kernel, centroid_index=1)
-
+    original_layer.build(input_shape)
+    model = keras.Sequential([original_layer])
+
+    # Save and load the layer in a temp directory
+    with tempfile.TemporaryDirectory() as tmp_dir_name:
+      keras_file = os.path.join(tmp_dir_name, 'keras_model')
+      keras.models.save_model(model, keras_file)
+      with cluster.cluster_scope():
+        loaded_layer = keras.models.load_model(keras_file).layers[0]
+
+    def assert_list_of_variables_all_equal(l1, l2):
+      self.assertLen(
+          l1, len(l2),
+          'lists l1 and l2 are not equal: \n l1={l1} \n l2={l2}'.format(
+              l1=[v.name for v in l1],
+              l2=[v.name for v in l2]))
+
+      name_to_var_from_l1 = {var.name: var for var in l1}
+      for var2 in l2:
+        self.assertIn(var2.name, name_to_var_from_l1)
+        arr1 = name_to_var_from_l1[var2.name].numpy()
+        arr2 = var2.numpy()
+        self.assertAllEqual(arr1, arr2)
+
+    # Check that trainable_weights and non_trainable_weights are the same
+    # in the original layer and loaded layer
+    assert_list_of_variables_all_equal(original_layer.trainable_weights,
+                                       loaded_layer.trainable_weights)
+    assert_list_of_variables_all_equal(original_layer.non_trainable_weights,
+                                       loaded_layer.non_trainable_weights)
 
 if __name__ == '__main__':
   test.main()