Implement update_training_weight functionality and add an example for it.

PiperOrigin-RevId: 368583345

Author: Xhark

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

@@ -37,6 +37,7 @@ class WeightCompressor(metaclass=abc.ABCMeta):
 
      This interface is a purely functional one.
   """
+  update_ops = []  # type: List
 
   # TODO(tfmot): Consider separate from algorithm API for custom layer supports.
   def get_compressible_weights(
@@ -103,19 +104,40 @@ def project_training_weights(
        tf.Tensor to set the compressible weight to.
     """
 
+  def init_update_ops(self, tensor_weight_pairs):
+    self.update_ops = []
+    self.tensor_weight_pairs = tensor_weight_pairs
+
   def update_training_weight(
-      self, training_weight: tf.Tensor, tensor: tf.Tensor):
-    """Update a training weight to a given tensor value.
+      self, training_weight: tf.Tensor, value: tf.Tensor):
+    """Add training weight assign op to the model update list.
+
+     This method is for the case that training weight should update to a
+    specific value not from the model optimizer. It will throw an error if it
+    can't find the training weight.
 
-    This method is for the case that training weight should update to a specific
-    value not from the model optimizer. It will throw an error if it can't
-    find the training weight.
+     This method should called in project_training_weights. During the training,
+    We collect all update_training_weight calls and make an UpdateOp for each
+    call. Finally, we put all these update ops to model.add_update.
 
     Args:
       training_weight: tf.Tensor representing a training weight.
-      tensor: tf.Tensor representing a value to be assigned to the training
+      value: tf.Tensor representing a value to be assigned to the training
         weight.
+    Raises:
+      ValueError if it can't find the training weight.
     """
+    for tensor, weight in self.tensor_weight_pairs:
+      if training_weight is tensor:
+        self.update_ops.append(weight.assign(value))
+        return
+
+    raise ValueError('Training weight not found. Please call '
+                     'the update_training_weight with given training '
+                     'weight tensor.')
+
+  def get_update_ops(self):
+    return self.update_ops
 
   def compress_training_weights(
       self, *training_weights: tf.Tensor) -> List[tf.Tensor]:

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

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

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

@@ -0,0 +1,166 @@
+# 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.
+# ==============================================================================
+"""SVD algorithm, where the training and inference graphs are the same."""
+from typing import List
+
+import tensorflow as tf
+
+from tensorflow_model_optimization.python.core.common.keras.compression import algorithm
+
+
+class SVD(algorithm.WeightCompressor):
+  """Define how to apply SVD algorithm.
+
+  This periodic update and scheduling base SVD algorithm update the original
+  weights to make lower rank by SVD for each update_freq steps. During the
+  warmup steps, It adjust the rank from the original to target rank gradually.
+  """
+
+  def __init__(self, rank, update_freq=100, warmup_step=1000):
+    self.rank = rank
+    self.update_freq = update_freq
+    self.warmup_step = warmup_step
+
+  # 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(
+      self, pretrained_weight: tf.Tensor):
+    self.add_training_weight(
+        name='w',
+        shape=pretrained_weight.shape,
+        dtype=pretrained_weight.dtype,
+        initializer=tf.keras.initializers.Constant(pretrained_weight))
+    self.add_training_weight(
+        name='step',
+        shape=(),
+        dtype=tf.int32,
+        initializer=tf.keras.initializers.Constant(0))
+
+  def decompress_weights(self, u: tf.Tensor, sv: tf.Tensor) -> tf.Tensor:
+    return tf.matmul(u, sv)
+
+  def project_training_weights(
+      self, weight: tf.Tensor, step: tf.Tensor) -> tf.Tensor:
+    weight_rank = tf.math.minimum(weight.shape[-1], weight.shape[-2])
+    self.update_training_weight(step, step + 1)
+    if step % self.update_freq == 0:
+      rank = self.rank
+      if step < self.warmup_step:
+        rank = tf.cast(tf.math.round(
+            weight_rank * (self.warmup_step - step)
+            + self.rank * step
+            ) / self.warmup_step, tf.int32)
+      rank = tf.math.minimum(rank, weight_rank)
+
+      s, u, v = tf.linalg.svd(weight)
+
+      if len(weight.shape) == 2:
+        # FC Layer
+        s = s[:rank]
+        u = u[:, :rank]
+        v = v[:, :rank]
+      elif len(weight.shape) == 4:
+        # Conv2D Layer
+        s = s[:, :, :rank]
+        u = u[:, :, :, :rank]
+        v = v[:, :, :, :rank]
+      else:
+        raise NotImplementedError('Only for dimension=2 or 4 is supported.')
+
+      sv = tf.matmul(tf.linalg.diag(s), v, adjoint_b=True)
+
+      new_weight = tf.matmul(u, sv)
+      self.update_training_weight(weight, new_weight)
+
+    return weight
+
+  def compress_training_weights(self, weight: tf.Tensor, _) -> List[tf.Tensor]:
+    rank = self.rank
+    s, u, v = tf.linalg.svd(weight)
+
+    if len(weight.shape) == 2:
+      # FC Layer
+      s = s[:rank]
+      u = u[:, :rank]
+      v = v[:, :rank]
+    elif len(weight.shape) == 4:
+      # Conv2D Layer
+      s = s[:, :, :rank]
+      u = u[:, :, :, :rank]
+      v = v[:, :, :, :rank]
+    else:
+      raise NotImplementedError('Only for dimension=2 or 4 is supported.')
+
+    sv = tf.matmul(tf.linalg.diag(s), v, adjoint_b=True)
+
+    return [u, sv]
+
+  def get_compressible_weights(
+      self, original_layer: tf.keras.layers.Layer) -> List[str]:
+    if isinstance(original_layer, (tf.keras.layers.Conv2D,
+                                   tf.keras.layers.Dense)):
+      return [original_layer.kernel]
+    return []
+
+  def optimize_model(self, to_optimize: tf.keras.Model) -> tf.keras.Model:
+    """Model developer API for optimizing a model for training.
+
+    The returned model should be used for compression aware training.
+    Args:
+      to_optimize: The model to be optimize.
+    Returns:
+      A wrapped model that has compression optimizers.
+    """
+    # pylint: disable=protected-access
+    if not isinstance(
+        to_optimize, tf.keras.Sequential) and not to_optimize._is_graph_network:
+      raise ValueError(
+          '`optimize_model` can only either be a tf.keras Sequential or '
+          'Functional model.')
+    # pylint: enable=protected-access
+
+    def _optimize_layer(layer):
+      # Require layer to be built so that the SVD-factorized weights
+      # can be initialized from the weights.
+      if not layer.built:
+        raise ValueError(
+            'Applying SVD currently requires passing in a built model')
+
+      return algorithm.create_layer_for_training(layer, algorithm=self)
+
+    return tf.keras.models.clone_model(
+        to_optimize, clone_function=_optimize_layer)
+
+  def compress_model(self, to_compress: tf.keras.Model) -> tf.keras.Model:
+    """Model developer API for optimizing a model for inference.
+
+    Args:
+      to_compress: The model that trained for compression. This model should
+        generated from the `optimize_model` method.
+    Returns:
+      A compressed model for the inference.
+    """
+    def _optimize_layer(layer):
+      # Require layer to be built so that the SVD-factorized weights
+      # can be initialized from the weights.
+      if not layer.built:
+        raise ValueError(
+            'Applying SVD currently requires passing in a built model')
+
+      return algorithm.create_layer_for_inference(layer, algorithm=self)
+
+    return tf.keras.models.clone_model(
+        to_compress, clone_function=_optimize_layer)