Merge branch 'master' into clusterable_layer

Author: wwwind

RELEASE.md

@@ -30,6 +30,14 @@ Keras pruning API:
 * Tested against TensorFlow 1.14.0, 2.0.0, and nightly, and Python 3.
 
 
+# TensorFlow Model Optimization next release TBD
+
+Keras clustering API:
+
+* Added *ClusteringSummaries* to create additional output for the clustering
+progress for TensorBoard.
+* Tested against TensorFlow 1.14.0, 2.0.0, and nightly, and Python 3.
+
 # TensorFlow Model Optimization 0.5.0
 
 TFMOT 0.5.0 adds some additional features for Quantization Aware Training. QAT

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

@@ -23,4 +23,6 @@
 
 from tensorflow_model_optimization.python.core.clustering.keras.cluster_config import CentroidInitialization
 from tensorflow_model_optimization.python.core.clustering.keras.clustering_algorithm import AbstractClusteringAlgorithm
+from tensorflow_model_optimization.python.core.clustering.keras.clustering_callbacks import ClusteringSummaries
+
 # pylint: enable=g-bad-import-order

tensorflow_model_optimization/python/core/clustering/keras/BUILD

@@ -12,6 +12,7 @@ py_library(
     srcs_version = "PY3",
     deps = [
         ":cluster",
+        ":clustering_callbacks",
         "//tensorflow_model_optimization/python/core/clustering/keras/experimental",
     ],
 )

tensorflow_model_optimization/python/core/internal/tensor_encoding/utils/tf_utils.py

@@ -311,6 +311,10 @@ def pack_into_int(value, input_bitrange, target_bitrange):
   """
   # TODO(b/161433177): Provide a general solution without memory overhead.
   # Special cases implemented without extra memory overhead.
+  if input_bitrange == 6 and target_bitrange == 28:
+    return _pack_into_int_6_28(value)
+  if input_bitrange == 7 and target_bitrange == 28:
+    return _pack_into_int_7_28(value)
   if input_bitrange == 8 and target_bitrange == 28:
     return _pack_into_int_8_28(value)
   if input_bitrange == 12 and target_bitrange == 28:
@@ -349,6 +353,10 @@ def unpack_from_int(value, original_bitrange, target_bitrange, shape):
   """
   # TODO(b/161433177): Provide a general solution without memory overhead.
   # Special cases implemented without extra memory overhead.
+  if original_bitrange == 6 and target_bitrange == 28:
+    return _unpack_from_int_6_28(value, shape)
+  if original_bitrange == 7 and target_bitrange == 28:
+    return _unpack_from_int_7_28(value, shape)
   if original_bitrange == 8 and target_bitrange == 28:
     return _unpack_from_int_8_28(value, shape)
   if original_bitrange == 12 and target_bitrange == 28:
@@ -386,6 +394,85 @@ def _expand_to_binary_form(value, input_bits):
   return tf.reshape(bits, [-1])
 
 
+def _pack_into_int_6_28(value):
+  """Implementation of `pack_into_int` for specific bitranges.
+
+  This method corresponts to `(input_bitrange, target_bitrange)` form the
+  `pack_into_int` method equal to `(6, 28)`. This method relies on the fact that
+  14 values in 6-bit bitrange can be packed into 3 values in 28-bitrange
+  (14 = least_common_multiple(6, 28) / 6).
+
+  It reshapes the input into matrix of 14 columns and performs operations on the
+  columns of the matrix, thus vectorizing the operations and avoiding memory
+  overhead of an earlier general implementation.
+
+  Args:
+    value: An integer Tensor to be packed with values in [0, 2**6 - 1].
+
+  Returns:
+    An integer Tensor representing `value` of the same dtype as `value`.
+  """
+  value = tf.reshape(value, [-1])
+  extra_zeros = tf.zeros(tf.math.mod(-tf.shape(value), 14), value.dtype)
+  val = tf.reshape(tf.concat([value, extra_zeros], 0), [-1, 14])
+
+  a = (val[:, 0] +
+       val[:, 1] * 2**6 +
+       val[:, 2] * 2**12 +
+       val[:, 3] * 2**18 +
+       tf.math.mod(val[:, 4], 2**4) * 2**24)
+  b = (tf.math.floordiv(val[:, 4], 2**4) +
+       val[:, 5] * 2**2 +
+       val[:, 6] * 2**8 +
+       val[:, 7] * 2**14 +
+       val[:, 8] * 2**20 +
+       tf.math.mod(val[:, 9], 2**2) * 2**26)
+  c = (tf.math.floordiv(val[:, 9], 2**2) +
+       val[:, 10] * 2**4 +
+       val[:, 11] * 2**10 +
+       val[:, 12] * 2**16 +
+       val[:, 13] * 2**22)
+
+  packed_val = tf.reshape(tf.stack([a, b, c], 1), [-1, 1])
+  if extra_zeros.shape[0] in [5, 6, 7, 8, 9]:
+    # We added unnecessary product of zeros to the representation.
+    packed_val = tf.slice(packed_val, [0, 0], [packed_val.shape[0] - 1, 1])
+  if extra_zeros.shape[0] in [10, 11, 12, 13]:
+    # We added unnecessary two products of zeros to the representation.
+    packed_val = tf.slice(packed_val, [0, 0], [packed_val.shape[0] - 2, 1])
+  return packed_val
+
+
+def _pack_into_int_7_28(value):
+  """Implementation of `pack_into_int` for specific bitranges.
+
+  This method corresponts to `(input_bitrange, target_bitrange)` form the
+  `pack_into_int` method equal to `(7, 28)`. This method relies on the fact that
+  4 values in 7-bit bitrange can be packed into 1 value in 28-bitrange
+  (4 = least_common_multiple(7, 28) / 7).
+
+  It reshapes the input into matrix of 4 columns and performs operations on the
+  columns of the matrix, thus vectorizing the operations and avoiding memory
+  overhead of an earlier general implementation.
+
+  Args:
+    value: An integer Tensor to be packed with values in [0, 2**7 - 1].
+
+  Returns:
+    An integer Tensor representing `value` of the same dtype as `value`.
+  """
+  value = tf.reshape(value, [-1])
+  extra_zeros = tf.zeros(tf.math.mod(-tf.shape(value), 4), value.dtype)
+  val = tf.reshape(tf.concat([value, extra_zeros], 0), [-1, 4])
+
+  packed_val = (val[:, 0] +
+                val[:, 1] * 2**7 +
+                val[:, 2] * 2**14 +
+                val[:, 3] * 2**21)
+
+  return tf.reshape(packed_val, [-1, 1])
+
+
 def _pack_into_int_8_28(value):
   """Implementation of `pack_into_int` for specific bitranges.
 
@@ -466,6 +553,47 @@ def _pack_into_int_12_28(value):
   return packed_val
 
 
+def _unpack_from_int_6_28(value, shape):
+  """Inverse operation of `_pack_into_int_6_28`."""
+  value = tf.reshape(value, [-1])
+  extra_zeros = tf.zeros(tf.math.mod(-tf.shape(value), 3), value.dtype)
+  val = tf.reshape(tf.concat([value, extra_zeros], 0), [-1, 3])
+
+  a = tf.math.mod(val[:, 0], 2**6)
+  b = tf.math.mod(tf.math.floordiv(val[:, 0], 2**6), 2**6)
+  c = tf.math.mod(tf.math.floordiv(val[:, 0], 2**12), 2**6)
+  d = tf.math.mod(tf.math.floordiv(val[:, 0], 2**18), 2**6)
+  e = tf.math.floordiv(val[:, 0], 2**24) + tf.math.mod(val[:, 1], 2**2) * 2**4
+  f = tf.math.mod(tf.math.floordiv(val[:, 1], 2**2), 2**6)
+  g = tf.math.mod(tf.math.floordiv(val[:, 1], 2**8), 2**6)
+  h = tf.math.mod(tf.math.floordiv(val[:, 1], 2**14), 2**6)
+  i = tf.math.mod(tf.math.floordiv(val[:, 1], 2**20), 2**6)
+  j = tf.math.floordiv(val[:, 1], 2**26) + tf.math.mod(val[:, 2], 2**4) * 2**2
+  k = tf.math.mod(tf.math.floordiv(val[:, 2], 2**4), 2**6)
+  l = tf.math.mod(tf.math.floordiv(val[:, 2], 2**10), 2**6)
+  m = tf.math.mod(tf.math.floordiv(val[:, 2], 2**16), 2**6)
+  n = tf.math.mod(tf.math.floordiv(val[:, 2], 2**22), 2**6)
+
+  unpacked_val = tf.reshape(
+      tf.stack([a, b, c, d, e, f, g, h, i, j, k, l, m, n], 1), [-1,])
+  unpacked_val = tf.slice(unpacked_val, [0], [tf.reduce_prod(shape)])
+  return tf.reshape(unpacked_val, shape)
+
+
+def _unpack_from_int_7_28(value, shape):
+  """Inverse operation of `_pack_into_int_7_28`."""
+  val = tf.reshape(value, [-1, 1])
+
+  a = tf.math.mod(val[:, 0], 2**7)
+  b = tf.math.mod(tf.math.floordiv(val[:, 0], 2**7), 2**7)
+  c = tf.math.mod(tf.math.floordiv(val[:, 0], 2**14), 2**7)
+  d = tf.math.mod(tf.math.floordiv(val[:, 0], 2**21), 2**7)
+
+  unpacked_val = tf.reshape(tf.stack([a, b, c, d], 1), [-1,])
+  unpacked_val = tf.slice(unpacked_val, [0], [tf.reduce_prod(shape)])
+  return tf.reshape(unpacked_val, shape)
+
+
 def _unpack_from_int_8_28(value, shape):
   """Inverse operation of `_pack_into_int_8_28`."""
   value = tf.reshape(value, [-1])

tensorflow_model_optimization/python/core/internal/tensor_encoding/utils/tf_utils_test.py

@@ -375,7 +375,8 @@ def test_unpack_from_int_different_outputs(self):
         packed_value, original_bitrange=4, target_bitrange=28, shape=(2,))
     self.assertAllEqual([9, 0], self.evaluate(unpacked_value))
 
-  @parameterized.parameters([(1, 28), (2, 28), (8, 28), (12, 28)])
+  @parameterized.parameters([(1, 28), (2, 28), (6, 28), (7, 28), (8, 28),
+                             (12, 28)])
   def test_boundary_conditions(self, input_bitrange, target_bitrange):
     max_v = 2**input_bitrange - 1
     input_value = tf.constant(
@@ -390,7 +391,8 @@ def test_boundary_conditions(self, input_bitrange, target_bitrange):
 
     self.assertAllEqual(self.evaluate(value), self.evaluate(unpacked_value))
 
-  @parameterized.parameters([(1, 28), (2, 28), (8, 28), (12, 28)])
+  @parameterized.parameters([(1, 28), (2, 28), (6, 28), (7, 28), (8, 28),
+                             (12, 28)])
   def test_random_input(self, input_bitrange, target_bitrange):
     # Tests that packing/unpacking amounts to identity, regardless of the input.
     num_elements = np.random.randint(low=1, high=50)
@@ -407,6 +409,42 @@ def test_random_input(self, input_bitrange, target_bitrange):
     except:  # pylint: disable=bare-except
       self.fail(f'Random input test failed with input value: {value}')
 
+  def test_pack_into_int_special_case_6_28(self):
+    value = tf.constant(
+        [50, 19, 51, 59, 10, 53, 36, 44, 31, 44, 31, 10, 31, 56, 49, 48, 35])
+    packed_value = tf_utils.pack_into_int(
+        value, input_bitrange=6, target_bitrange=28)
+    expected_packed_value = tf.constant([[183448818], [33236180], [236923387],
+                                         [146481]])
+    self.assertAllEqual(self.evaluate(expected_packed_value),
+                        self.evaluate(packed_value))
+
+  def test_unpack_from_int_special_case_6_28(self):
+    packed_value = tf.constant([[183448818], [33236180], [236923387], [146481]])
+    unpacked_value = tf_utils.unpack_from_int(
+        packed_value, original_bitrange=6, target_bitrange=28, shape=(17,))
+    expected_unpacked_value = tf.constant(
+        [50, 19, 51, 59, 10, 53, 36, 44, 31, 44, 31, 10, 31, 56, 49, 48, 35])
+    self.assertAllEqual(self.evaluate(expected_unpacked_value),
+                        self.evaluate(unpacked_value))
+
+  def test_pack_into_int_special_case_7_28(self):
+    value = tf.constant([117, 86, 42, 69, 9, 70, 66, 8, 112, 116])
+    packed_value = tf_utils.pack_into_int(
+        value, input_bitrange=7, target_bitrange=28)
+    expected_packed_value = tf.constant([[145402741], [17867529], [14960]])
+    self.assertAllEqual(self.evaluate(expected_packed_value),
+                        self.evaluate(packed_value))
+
+  def test_unpack_from_int_special_case_7_28(self):
+    packed_value = tf.constant([[145402741], [17867529], [14960]])
+    unpacked_value = tf_utils.unpack_from_int(
+        packed_value, original_bitrange=7, target_bitrange=28, shape=(10,))
+    expected_unpacked_value = tf.constant(
+        [117, 86, 42, 69, 9, 70, 66, 8, 112, 116])
+    self.assertAllEqual(self.evaluate(expected_unpacked_value),
+                        self.evaluate(unpacked_value))
+
   def test_pack_into_int_special_case_8_28(self):
     value = tf.constant([38, 147, 1, 201, 205, 36, 155, 78, 163, 98])
     packed_value = tf_utils.pack_into_int(

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

@@ -92,7 +92,6 @@ class Default8BitQuantizeRegistry(
       # layers.DepthwiseConv2D is supported and handled in code below.
 
       # _QuantizeInfo(layers.Conv3D, ['kernel'], ['activation']),
-      # _QuantizeInfo(layers.Conv2DTranspose, ['kernel'], ['activation']),
       # _QuantizeInfo(layers.Conv3DTranspose, ['kernel'], ['activation']),
       _no_quantize(layers.Cropping1D),
       _no_quantize(layers.Cropping2D),
@@ -198,6 +197,9 @@ def __init__(self):
     self._layer_quantize_map[
         layers.DepthwiseConv2D] = Default8BitConvQuantizeConfig(
             ['depthwise_kernel'], ['activation'], False)
+    self._layer_quantize_map[layers.Conv2DTranspose] = \
+        Default8BitConvTransposeQuantizeConfig(
+            ['kernel'], ['activation'], False)
 
   def _is_supported_layer(self, layer_class):
     return layer_class in self._layer_quantize_map
@@ -509,6 +511,17 @@ def __init__(self, weight_attrs, activation_attrs, quantize_output):
     )
 
 
+class Default8BitConvTransposeQuantizeConfig(Default8BitQuantizeConfig):
+  """QuantizeConfig for Conv2DTranspose layers."""
+
+  def __init__(self, weight_attrs, activation_attrs, quantize_output):
+    super(Default8BitConvTransposeQuantizeConfig,
+          self).__init__(weight_attrs, activation_attrs, quantize_output)
+
+    self.weight_quantizer = default_8bit_quantizers.Default8BitConvTransposeWeightsQuantizer(
+    )
+
+
 def _types_dict():
   return {
       'Default8BitQuantizeConfig':

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

@@ -41,3 +41,20 @@ def build(self, tensor_shape, name, layer):
         trainable=False)
 
     return {'min_var': min_weight, 'max_var': max_weight}
+
+
+class Default8BitConvTransposeWeightsQuantizer(quantizers.LastValueQuantizer):
+  """Quantizer for handling weights in Conv2DTranspose layers."""
+
+  def __init__(self):
+    """Construct LastValueQuantizer with params specific for TFLite Conv2DTranpose."""
+
+    super(Default8BitConvTransposeWeightsQuantizer, self).__init__(
+        num_bits=8, per_axis=False, symmetric=True, narrow_range=True)
+
+  def __call__(self, inputs, training, weights, **kwargs):
+    outputs = tf.transpose(inputs, (0, 1, 3, 2))
+    outputs = super(Default8BitConvTransposeWeightsQuantizer,
+                    self).__call__(outputs, training, weights, **kwargs)
+    outputs = tf.transpose(outputs, (0, 1, 3, 2))
+    return outputs

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

@@ -153,6 +153,13 @@ def _get_upsampling2d_bilinear_model(self):
     x = tf.keras.layers.UpSampling2D(size=(1, 5), interpolation='bilinear')(i)
     return tf.keras.Model(i, x)
 
+  def _get_conv2d_transpose_model(self):
+    i = tf.keras.Input(shape=(32, 32, 3))
+    x = tf.keras.layers.Conv2DTranspose(
+        2, kernel_size=(3, 3), strides=(2, 2))(
+            i)
+    return tf.keras.Model(i, x)
+
   @parameterized.parameters([
       _get_single_conv_model, _get_single_dense_model,
       _get_single_conv_relu_model, _get_stacked_convs_model,
@@ -165,6 +172,7 @@ def _get_upsampling2d_bilinear_model(self):
       # TODO(tfmot): There are gaps between ResizeBilinear with FakeQuant and
       # TFLite quantized ResizeBilinear op. It has a bit more quantization
       # error than other ops in this test now.
+      _get_conv2d_transpose_model,
   ])
   def testModelEndToEnd(self, model_fn):
     # 1. Check whether quantized model graph can be constructed.

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

@@ -298,7 +298,6 @@ class QuantizeFullIntegerModelTest(tf.test.TestCase, parameterized.TestCase):
           layers.UpSampling3D,
           # Not done since not registered since not per-axis yet.
           layers.Conv1D,
-          layers.Conv2DTranspose,
       ]
   ])
   def testQuantizeSingleLayer_ProducesFullIntegerModel_TF2(

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -91,6 +91,7 @@ py_library(
         # numpy dep1,
         # tensorflow dep1,
         # python/keras/utils:generic_utils tensorflow dep2,
+        "//tensorflow_model_optimization/python/core/keras:compat",
         "//tensorflow_model_optimization/python/core/keras:utils",
     ],
 )