Add monitoring for API usage

PiperOrigin-RevId: 365508776

Author: daverim

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_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')

tensorflow_model_optimization/python/core/quantization/keras/quantize.py

@@ -16,6 +16,7 @@
 
 import tensorflow as tf
 
+from tensorflow_model_optimization.python.core.keras import metrics
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_annotate as quantize_annotate_mod
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_aware_activation
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_config as quantize_config_mod
@@ -263,6 +264,7 @@ def quantize_annotate_layer(to_annotate, quantize_config=None):
       layer=to_annotate, quantize_config=quantize_config)
 
 
+@metrics.MonitorBoolGauge('quantize_apply_usage')
 def quantize_apply(
     model,
     scheme=default_8bit_quantize_scheme.Default8BitQuantizeScheme()):

tensorflow_model_optimization/python/core/quantization/keras/quantize_wrapper.py

@@ -30,6 +30,7 @@
 
 from tensorflow.python.util import tf_inspect
 
+from tensorflow_model_optimization.python.core.keras import metrics
 from tensorflow_model_optimization.python.core.keras import utils
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_aware_activation
 
@@ -70,6 +71,8 @@ def __init__(self, layer, quantize_config, **kwargs):
     self.quantize_config = quantize_config
 
     self._track_trackable(layer, name='layer')
+    metrics.MonitorBoolGauge('quantize_wrapper_usage').set(
+        layer.__class__.__name__)
 
   @staticmethod
   def _make_layer_name(layer):

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -27,6 +27,7 @@ py_library(
         ":pruning_schedule",
         ":pruning_wrapper",
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/keras:metrics",
     ],
 )
 
@@ -92,6 +93,7 @@ py_library(
         # tensorflow dep1,
         # python/keras/utils:generic_utils tensorflow dep2,
         "//tensorflow_model_optimization/python/core/keras:compat",
+        "//tensorflow_model_optimization/python/core/keras:metrics",
         "//tensorflow_model_optimization/python/core/keras:utils",
     ],
 )

tensorflow_model_optimization/python/core/sparsity/keras/prune.py

@@ -17,6 +17,7 @@
 
 import tensorflow as tf
 
+from tensorflow_model_optimization.python.core.keras import metrics
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule as pruning_sched
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper
 
@@ -52,6 +53,7 @@ def prune_scope():
       {'PruneLowMagnitude': pruning_wrapper.PruneLowMagnitude})
 
 
+@metrics.MonitorBoolGauge('prune_low_magnitude_usage')
 def prune_low_magnitude(to_prune,
                         pruning_schedule=pruning_sched.ConstantSparsity(0.5, 0),
                         block_size=(1, 1),

tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py

@@ -27,6 +27,7 @@
 # b/(139939526): update to use public API.
 from tensorflow.python.keras.utils import generic_utils
 from tensorflow_model_optimization.python.core.keras import compat as tf_compat
+from tensorflow_model_optimization.python.core.keras import metrics
 from tensorflow_model_optimization.python.core.keras import utils
 from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
 from tensorflow_model_optimization.python.core.sparsity.keras import prune_registry
@@ -185,6 +186,8 @@ def __init__(self,
     if not hasattr(self, '_batch_input_shape') and hasattr(
         layer, '_batch_input_shape'):
       self._batch_input_shape = self.layer._batch_input_shape
+    metrics.MonitorBoolGauge('prune_low_magnitude_wrapper_usage').set(
+        layer.__class__.__name__)
 
   def build(self, input_shape):
     super(PruneLowMagnitude, self).build(input_shape)