Initial API for specifying 'compressible_weights'.

For testing, Added bias_only algorithm that compress bias vector has same weight for each layer.

PiperOrigin-RevId: 339374777

Author: Xhark

tensorflow_model_optimization/python/core/common/keras/compression/algorithm.py

@@ -110,6 +110,23 @@ def training(self, *training_weights: tf.Tensor) -> tf.Tensor:
        tf.Tensor to set the compressible weight to.
     """
 
+  # TODO(tfmot): Consider separate from algorithm API for custom layer supports.
+  def get_compressible_weights(
+      self, original_layer: tf.keras.layers.Layer) -> List[str]:
+    """Define compressible weights for each layer.
+
+    Args:
+       original_layer: tf.keras.layers.Layer representing a layer from the
+       original model.
+
+    Returns:
+       List of atrribute names as string representing list of compressible
+       weights for the given layer. (e.g. return value ['kernel'] means
+       layer.kernel is compressible.)
+    """
+    del original_layer
+    return []
+
 
 def create_layer_for_training(
     layer: tf.keras.layers.Layer,

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

@@ -45,3 +45,25 @@ py_test(
         # tensorflow dep1,
     ],
 )
+
+py_library(
+    name = "bias_only",
+    srcs = ["bias_only.py"],
+    srcs_version = "PY3",
+    deps = [
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/common/keras/compression:algorithm",
+    ],
+)
+
+py_test(
+    name = "bias_only_test",
+    timeout = "long",
+    srcs = ["bias_only_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":bias_only",
+        # numpy dep1,
+        # tensorflow dep1,
+    ],
+)

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

@@ -0,0 +1,80 @@
+# 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.
+# ==============================================================================
+"""BiasOnly algorithm, where the compress bias only."""
+from typing import List
+
+import tensorflow as tf
+
+from tensorflow_model_optimization.python.core.common.keras.compression import algorithm
+
+
+# TODO(tfmot): This algorithm is showcase for bias only compression. if we find
+# better algorithm that can show better compressible weights coverage, then
+# we can remove this algorithm.
+class BiasOnly(algorithm.WeightCompressionAlgorithm):
+  """Define how to apply BiasOnly algorithm."""
+
+  # TODO(tfmot): communicate that `pretrained_weight` will sometimes
+  # be a dummy tensor and sometimes be actual pretrained values during
+  # its actual usage.
+  def init_training_weights_repr(
+      self, pretrained_weight: tf.Tensor) -> List[algorithm.WeightRepr]:
+    bias_mean = tf.reduce_mean(pretrained_weight)
+    bias_shape = tf.shape(pretrained_weight)
+
+    # TODO(tfmot): note that it does not suffice to just have the initializer
+    # to derive the shape from, in the case of a constant initializer.
+    # The unit test fail without providing the shape.
+    return [
+        algorithm.WeightRepr(
+            name='bias_mean',
+            shape=(),
+            initializer=tf.keras.initializers.Constant(bias_mean)),
+        algorithm.WeightRepr(
+            name='bias_shape',
+            shape=bias_shape.shape,
+            dtype=bias_shape.dtype,
+            initializer=tf.keras.initializers.Constant(bias_shape))
+    ]
+
+  def decompress(
+      self, bias_mean: tf.Tensor, bias_shape: tf.Tensor) -> tf.Tensor:
+    return tf.broadcast_to(bias_mean, bias_shape)
+
+  def training(
+      self, bias_mean: tf.Tensor, bias_shape: tf.Tensor) -> tf.Tensor:
+    return self.decompress(bias_mean, bias_shape)
+
+  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 ['bias']
+    return []
+
+
+def optimize(to_optimize: tf.keras.Model) -> tf.keras.Model:
+  """Model developer API for optimizing a model."""
+
+  def _optimize_layer(layer):
+    # Require layer to be built so that the average of bias can be initialized.
+    if not layer.built:
+      raise ValueError(
+          'Applying BiasOnly currently requires passing in a built model')
+
+    return algorithm.create_layer_for_training(layer, algorithm=BiasOnly())
+
+  return tf.keras.models.clone_model(
+      to_optimize, clone_function=_optimize_layer)

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

@@ -0,0 +1,217 @@
+# 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 bias only optimization."""
+
+import os
+import tempfile
+
+import numpy as np
+import tensorflow as tf
+
+from tensorflow_model_optimization.python.core.common.keras.compression.algorithms import bias_only
+
+
+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[0:1000]
+  y_train = y_train[0: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
+
+
+# TODO(tfmot): reuse existing test utilities.
+def _convert_to_tflite(saved_model_dir):
+  _, tflite_file = tempfile.mkstemp()
+
+  converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
+  tflite_model = converter.convert()
+
+  with open(tflite_file, 'wb') as f:
+    f.write(tflite_model)
+
+  return tflite_file
+
+
+# TODO(tfmot): reuse test_utils_mnist.py.
+def _test_tflite(tflite_file):
+  interpreter = tf.lite.Interpreter(model_path=tflite_file)
+  interpreter.allocate_tensors()
+
+  input_index = interpreter.get_input_details()[0]['index']
+  output_index = interpreter.get_output_details()[0]['index']
+
+  (_, _), (x_test, y_test) = _get_dataset()
+
+  # Testing the entire dataset is too slow. Verifying only 300 of 10k samples.
+  x_test = x_test[0:300, :]
+  y_test = y_test[0:300]
+
+  total_seen = 0
+  num_correct = 0
+
+  for img, label in zip(x_test, y_test):
+    batch_input_shape = (1, 28, 28)
+    inp = img.reshape(batch_input_shape)
+    inp = inp.astype(np.float32)
+    total_seen += 1
+    interpreter.set_tensor(input_index, inp)
+    interpreter.invoke()
+    predictions = interpreter.get_tensor(output_index)
+
+    if np.argmax(predictions) == label:
+      num_correct += 1
+
+  return float(num_correct) / float(total_seen)
+
+
+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 testBiasOnly_ReducesParamaters(self):
+    model = _build_model()
+    compressed_model = bias_only.optimize(model)
+
+    self.assertEqual(model.count_params(), 431080)
+    self.assertEqual(compressed_model.count_params(), 430508)
+
+  def testBiasOnly_HasReasonableAccuracy_TF(self):
+    model = _build_model()
+
+    compressed_model = bias_only.optimize(model)
+
+    _train_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)
+
+    self.assertGreater(results[1], 0.60)
+
+  def testBiasOnly_HasReasonableAccuracy_TFLite(self):
+    model = _build_model()
+
+    compressed_model = bias_only.optimize(model)
+
+    _train_model(compressed_model)
+
+    saved_model_dir = _save_as_saved_model(compressed_model)
+    compressed_tflite_file = _convert_to_tflite(saved_model_dir)
+
+    accuracy = _test_tflite(compressed_tflite_file)
+
+    self.assertGreater(accuracy, 0.60)
+
+  # TODO(tfmot): can simplify to single layer test.
+  def testBiasOnly_BreaksDownLayerWeights(self):
+    model = _build_model()
+
+    first_conv_layer = model.layers[2]
+    self.assertLen(first_conv_layer.weights, 2)
+
+    compressed_model = bias_only.optimize(model)
+
+    first_conv_layer = compressed_model.layers[2]
+
+    self.assertLen(first_conv_layer.weights, 3)
+
+  # TODO(tfmot): can simplify to single layer test.
+  def testBiasOnly_PreservesPretrainedWeights(self):
+    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()
+
+    compressed_model = bias_only.optimize(model)
+
+    dense_layer_compressed_weights = compressed_model.layers[1].get_weights()
+
+    # kernel
+    assert (dense_layer_weights[0] == dense_layer_compressed_weights[2]).all()
+
+    # bias
+    algorithm = bias_only.BiasOnly()
+    w1_repr, w2_repr = algorithm.init_training_weights_repr(
+        dense_layer_weights[1])
+
+    w1 = w1_repr.initializer(shape=None, dtype=w1_repr.dtype)
+    w2 = w2_repr.initializer(shape=None, dtype=w2_repr.dtype)
+
+    assert (w1 == dense_layer_compressed_weights[0]).numpy().all()
+    assert (w2 == dense_layer_compressed_weights[1]).numpy().all()
+
+
+if __name__ == '__main__':
+  tf.test.main()

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

@@ -73,6 +73,13 @@ def compress(self, weight: tf.Tensor) -> List[tf.Tensor]:
   def training(self, weight: tf.Tensor) -> tf.Tensor:
     return weight
 
+  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 []
+
 
 # TODO(tfmot): consider if we can simplify `create_model_for_training` and
 # `create_model_for_inference` into a single API for algorithm developers.

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

@@ -76,6 +76,13 @@ def decompress(self, u: tf.Tensor, sv: tf.Tensor) -> tf.Tensor:
   def training(self, u: tf.Tensor, sv: tf.Tensor) -> tf.Tensor:
     return self.decompress(u, sv)
 
+  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: SVDParams) -> tf.keras.Model:
   """Model developer API for optimizing a model."""