Merge branch 'master' into clusterable_layer

Author: wwwind

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

@@ -54,6 +54,25 @@ 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,
+    )
+    self.b = self.add_weight(
+      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):
@@ -108,6 +127,8 @@ def setUp(self):
     self.custom_non_clusterable_layer = CustomNonClusterableLayer(10)
     self.keras_depthwiseconv2d_layer = layers.DepthwiseConv2D((3, 3), (1, 1))
     self.clusterable_layer = MyClusterableLayer(10)
+    self.keras_custom_layer = KerasCustomLayer()
+
 
     clustering_registry.ClusteringLookupRegistry.register_new_implementation(
         {
@@ -252,6 +273,22 @@ def testClusterMyClusterableLayerInvalid(self):
     with self.assertRaises(TypeError):
       MyClusterableLayerInvalid(10) # pylint: disable=abstract-class-instantiated
 
+  def testClusterKerasCustomLayer(self):
+    """
+    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)
+    # 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)
+
   @keras_parameterized.run_all_keras_modes
   def testClusterSequentialModelSelectively(self):
     clustered_model = keras.Sequential()

tensorflow_model_optimization/python/core/keras/BUILD

@@ -43,3 +43,23 @@ py_library(
         # python:variables tensorflow dep2,
     ],
 )
+
+py_library(
+    name = "metrics",
+    srcs = ["metrics.py"],
+    srcs_version = "PY3",
+    deps = [
+        # python/eager:monitoring tensorflow dep2,
+    ],
+)
+
+py_test(
+    name = "metrics_test",
+    srcs = ["metrics_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":metrics",
+        # mock dep1,
+        # tensorflow dep1,
+    ],
+)

tensorflow_model_optimization/python/core/keras/metrics.py

@@ -0,0 +1,79 @@
+# 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.
+# ==============================================================================
+"""Implements monitoring."""
+
+from tensorflow.python.eager import monitoring
+
+
+class MonitorBoolGauge():
+  """Monitoring utility class for usage metrics."""
+
+  _PRUNE_LOW_MAGNITUDE_USAGE = monitoring.BoolGauge(
+      '/tfmot/api/sparsity/prune_low_magnitude',
+      'prune_low_magnitude usage.', 'status')
+
+  _PRUNE_WRAPPER_USAGE = monitoring.BoolGauge(
+      '/tfmot/api/sparsity/pruning_wrapper',
+      'Pruning wrapper class usage.', 'layer')
+
+  _QUANTIZE_APPLY_USAGE = monitoring.BoolGauge(
+      '/tfmot/api/quantization/quantize_apply',
+      'quantize_apply usage.', 'status')
+
+  _QUANTIZE_WRAPPER_USAGE = monitoring.BoolGauge(
+      '/tfmot/api/quantization/quantize_wrapper',
+      'Quantize wrapper class usage.', 'layer')
+
+  _SUCCESS_LABEL = 'success'
+  _FAILURE_LABEL = 'failure'
+
+  def __init__(self, name):
+    self.bool_gauge = self.get_usage_gauge(name)
+
+  def get_usage_gauge(self, name):
+    if name == 'prune_low_magnitude_usage':
+      return MonitorBoolGauge._PRUNE_LOW_MAGNITUDE_USAGE
+    if name == 'prune_low_magnitude_wrapper_usage':
+      return MonitorBoolGauge._PRUNE_WRAPPER_USAGE
+    if name == 'quantize_apply_usage':
+      return MonitorBoolGauge._QUANTIZE_APPLY_USAGE
+    if name == 'quantize_wrapper_usage':
+      return MonitorBoolGauge._QUANTIZE_WRAPPER_USAGE
+    raise ValueError('Invalid gauge name: {}'.format(name))
+
+  def __call__(self, func):
+    def inner(*args, **kwargs):
+      try:
+        results = func(*args, **kwargs)
+        self.bool_gauge.get_cell(MonitorBoolGauge._SUCCESS_LABEL).set(True)
+        return results
+      except Exception as error:
+        self.bool_gauge.get_cell(MonitorBoolGauge._FAILURE_LABEL).set(True)
+        raise error
+
+    if self.bool_gauge:
+      return inner
+
+    return func
+
+  def set(self, label=None, value=True):
+    """Set the bool gauge to value if initialized.
+
+    Args:
+      label: optional string label defaults to None.
+      value: optional bool value defaults to True.
+    """
+    if self.bool_gauge:
+      self.bool_gauge.get_cell(label).set(value)

tensorflow_model_optimization/python/core/keras/metrics_test.py

@@ -0,0 +1,73 @@
+# 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.
+# ==============================================================================
+"""Tests for Metrics."""
+
+import mock
+import tensorflow as tf
+
+from tensorflow.python.eager import monitoring
+from tensorflow_model_optimization.python.core.keras import metrics
+
+
+class MetricsTest(tf.test.TestCase):
+
+  gauge = monitoring.BoolGauge('/tfmot/metrics/testing', 'testing', 'labels')
+
+  def setUp(self):
+    super(MetricsTest, self).setUp()
+    self.test_label = tf.keras.layers.Conv2D(1, 1).__class__.__name__
+    for label in [
+        self.test_label, metrics.MonitorBoolGauge._SUCCESS_LABEL,
+        metrics.MonitorBoolGauge._FAILURE_LABEL
+    ]:
+      MetricsTest.gauge.get_cell(label).set(False)
+
+    with mock.patch.object(metrics.MonitorBoolGauge, 'get_usage_gauge',
+                           return_value=MetricsTest.gauge):
+      self.monitor = metrics.MonitorBoolGauge('testing')
+
+  def test_DecoratorTest(self):
+    @self.monitor
+    def func(x):
+      return x + 1
+
+    self.assertEqual(func(1), 2)
+    self.assertTrue(MetricsTest.gauge.get_cell(
+        metrics.MonitorBoolGauge._SUCCESS_LABEL).value())
+
+  def test_DecoratorFailureTest(self):
+    @self.monitor
+    def func(x):
+      raise ValueError()
+
+    with self.assertRaises(ValueError):
+      func(1)
+    self.assertTrue(MetricsTest.gauge.get_cell(
+        metrics.MonitorBoolGauge._FAILURE_LABEL).value())
+
+  def test_UndecoratedTest(self):
+    with self.assertRaises(ValueError):
+      @metrics.MonitorBoolGauge('unknown')
+      def func(x):
+        return x+1
+      func(1)
+
+  def test_SetTest(self):
+    self.monitor.set(self.test_label)
+    self.assertTrue(MetricsTest.gauge.get_cell(self.test_label).value())
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow_model_optimization/python/core/quantization/keras/BUILD

@@ -211,6 +211,7 @@ py_library(
         ":quantize_config",
         ":quantizers",
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/keras:metrics",
         "//tensorflow_model_optimization/python/core/keras:utils",
     ],
 )
@@ -244,6 +245,7 @@ py_library(
         ":quantize_layer",
         ":quantize_wrapper",
         # tensorflow dep1,
+        "//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/layers:conv_batchnorm",

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

@@ -65,8 +65,8 @@ def apply(self, model, layer_quantize_map):
         default_8bit_transforms.ConcatTransform4Inputs(),
         default_8bit_transforms.ConcatTransform3Inputs(),
         default_8bit_transforms.ConcatTransform(),
-        default_8bit_transforms.AddReLUQuantize(),
-        default_8bit_transforms.AddActivationQuantize(),
+        default_8bit_transforms.LayerReLUQuantize(),
+        default_8bit_transforms.LayerReluActivationQuantize(),
     ]
     return model_transformer.ModelTransformer(
         model, transforms,

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

@@ -497,11 +497,12 @@ def replacement(self, match_layer):
         metadata=conv_metadata)
 
 
-class AddReLUQuantize(transforms.Transform):
+class LayerReLUQuantize(transforms.Transform):
   """Ensure FQ does not get placed between Add and ReLU."""
 
   def pattern(self):
-    return LayerPattern('ReLU', inputs=[LayerPattern('Add')])
+    return LayerPattern(
+        'ReLU', inputs=[LayerPattern('Add|Conv2D|DepthwiseConv2D|Dense')])
 
   def replacement(self, match_layer):
     relu_layer_node = match_layer
@@ -518,14 +519,14 @@ def custom_objects(self):
     }
 
 
-class AddActivationQuantize(AddReLUQuantize):
+class LayerReluActivationQuantize(LayerReLUQuantize):
   """Ensure FQ does not get placed between Add and ReLU."""
 
   def pattern(self):
     return LayerPattern(
         'Activation',
         config={'activation': 'relu'},
-        inputs=[LayerPattern('Add')])
+        inputs=[LayerPattern('Add|Conv2D|DepthwiseConv2D|Dense')])
 
 
 class InputLayerQuantize(transforms.Transform):

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

@@ -45,6 +45,11 @@
 # TODO(alanchiao): reduce redundancy by parameterizing on Depthwise vs Conv.
 class DefaultTransformsTest(tf.test.TestCase, parameterized.TestCase):
 
+  @classmethod
+  def setUpClass(cls):
+    super(DefaultTransformsTest, cls).setUpClass()
+    np.random.seed(12345678)
+
   def testTransformsConvBNReLUPattern(self):
     model = Conv2DModel.get_nonfolded_batchnorm_model(
         post_bn_activation=keras.layers.ReLU(6.0), model_type='functional')
@@ -344,8 +349,8 @@ def testSeparableConvQuantize_(self, kwargs):
   # Conv2DReshapeBatchNormActivationQuantize
 
   @parameterized.parameters(
-      ('relu', default_8bit_transforms.AddReLUQuantize),
-      ('act_relu', default_8bit_transforms.AddActivationQuantize),
+      ('relu', default_8bit_transforms.LayerReLUQuantize),
+      ('act_relu', default_8bit_transforms.LayerReluActivationQuantize),
   )
   def testAddReLUQuantize(self, activation_type, transform_type):
     add = keras.layers.Add()
@@ -370,6 +375,33 @@ def testAddReLUQuantize(self, activation_type, transform_type):
         updated_metadata.get(add_layer.name).get('quantize_config'),
         default_8bit_quantize_configs.NoOpQuantizeConfig)
 
+  @parameterized.parameters(
+      ('relu', default_8bit_transforms.LayerReLUQuantize),
+      ('act_relu', default_8bit_transforms.LayerReluActivationQuantize))
+  def testLayerReLUQuantize(self, activation_type, transform_type):
+    # TODO(tfmot): Add tests for DepthConv and Dense
+    input_shape = (1, 3, 3, 3)
+    conv_layer = tf.keras.layers.Conv2D(5, 2, input_shape=input_shape)
+    if activation_type == 'relu':
+      act_layer = keras.layers.ReLU(6.0)
+    elif activation_type == 'act_relu':
+      act_layer = keras.layers.Activation('relu')
+
+    model = tf.keras.Sequential([conv_layer, act_layer])
+
+    transformed_model, updated_metadata = ModelTransformer(
+        model,
+        [transform_type()],
+    ).transform()
+
+    self.assertIsInstance(
+        updated_metadata.get(model.layers[0].name).get('quantize_config'),
+        default_8bit_quantize_configs.NoOpQuantizeConfig)
+
+    inputs = np.random.standard_normal(input_shape)
+    self.assertAllClose(
+        transformed_model.predict(inputs), model.predict(inputs))
+
   def testAddsQuantizeLayerAfterInputLayer(self):
     inp1 = keras.layers.Input((3,))
     inp2 = keras.layers.Input((3,))