Update compression API core to match with the RFC (tensorflow/community#342).

PiperOrigin-RevId: 361948873

Author: Xhark

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

@@ -24,21 +24,11 @@
 
 @dataclasses.dataclass
 class WeightRepr:
-  """Dataclass that wraps `tf.keras.layers.Layer.add_weight` parameters."""
-  name: Any = None
-  shape: Any = None
-  dtype: Any = None
-  initializer: Any = None
-  regularizer: Any = None
-  trainable: Any = None
-  constraint: Any = None
-  partitioner: Any = None
-  use_resource: Any = None
-  synchronization: Any = tf.VariableSynchronization.AUTO
-  aggregation: Any = tf.compat.v1.VariableAggregation.NONE
-
-
-class WeightCompressionAlgorithm(metaclass=abc.ABCMeta):
+  args: Any = None
+  kwargs: Any = None
+
+
+class WeightCompressor(metaclass=abc.ABCMeta):
   """Interface for weight compression algorithm that acts on a per-layer basis.
 
      This allows both options of either decompressing during inference or
@@ -48,93 +38,127 @@ class WeightCompressionAlgorithm(metaclass=abc.ABCMeta):
      This interface is a purely functional one.
   """
 
-  @abc.abstractmethod
-  def init_training_weights_repr(
-      self, pretrained_weight: tf.Tensor) -> List[WeightRepr]:
-    """Create training weight representations for initializing layer variables.
+  # TODO(tfmot): Consider separate from algorithm API for custom layer supports.
+  def get_compressible_weights(
+      self, original_layer: tf.keras.layers.Layer) -> List[tf.Variable]:
+    """Define compressible weights for each layer.
 
     Args:
-      pretrained_weight: tf.Tensor of a pretrained weight of a layer that will
-        be compressed eventually.
+       original_layer: tf.keras.layers.Layer representing a layer from the
+       original model.
 
     Returns:
-      A list of `WeightRepr`, a container for arguments to
-      `tf.keras.layers.Layer.add_weight`for each tf.Variable to create.
+       List of compressible weights for the given layer.
     """
+    del original_layer
+    return []
 
-  def compress(self, *training_weights: tf.Tensor) -> List[tf.Tensor]:
-    """Define the operations to compress a single weight after training.
-
-    'Compress' can refer to making the weight more amenable to compression
-    or actually compress the weight.
+  @abc.abstractmethod
+  def init_training_weights(
+      self, pretrained_weight: tf.Tensor):
+    """Initialize training weights for the compressible weight.
 
-    The default is an identity.
+    It calls the `add_training_weight` to add a training weight for a given
+    `pretrained_weight`. A `pretrained_weight` can have multiple training
+    weights. We initialize the training weights for each compressible
+    weight by just calling this function for each.
 
     Args:
-      *training_weights: tf.Tensors representing all variables used during
-        training, for a single compressible weight, in the order returned in
-        `init_training_weights_repr`.
-
-    Returns:
-      List of tf.Tensors to set to compressed or more compressible form.
+      pretrained_weight: tf.Tensor of a pretrained weight of a layer that will
+        be compressed eventually.
     """
-    return list(training_weights)
 
-  def decompress(self, *compressed_weights: tf.Tensor) -> tf.Tensor:
-    """Define the operations to decompress a single weight’s compressed form during inference.
+  def add_training_weight(
+      self, *args, **kwargs):
+    """Add a training weight for the compressible weight.
 
-    The default is an identity. TODO(): actually isn't.
+    When this method is called from the `init_training_weights`, this adds
+    training weights for the pretrained_weight that is the input of the
+    `init_training_weights`.
 
     Args:
-       *compressed_weights: tf.Tensors representing a single weight’s compressed
-         form, coming from what’s returned in `compress`.
-
-    Returns:
-      A tf.Tensor representing the decompressed `compressed_weights`.
+      *args: Passed through to training_model.add_weight.
+      **kwargs: Passed through to training_model.add_weight.
     """
-    return compressed_weights[0]
+    weight_repr = WeightRepr(args=args, kwargs=kwargs)
+    if hasattr(self, 'weight_reprs'):
+      self.weight_reprs.append(weight_repr)
+    else:
+      self.weight_reprs = [weight_repr]
 
   @abc.abstractmethod
-  def training(self, *training_weights: tf.Tensor) -> tf.Tensor:
-    """Define a piece of the forward pass during training, which operates on a single compressible weight.
+  def project_training_weights(
+      self, *training_weights: tf.Tensor) -> tf.Tensor:
+    """Define a piece of the forward pass during training.
 
-    TODO(tfmot): throw this error.
+    It operates on a single compressible weight.
     The default throws an error when training occurs.
 
     Args:
        *training_weights: tf.Tensors representing any variables used during
          training, for a single compressible weight, in the order returned in
-         `init_training_weights_repr`.
+         `init_training_weights`.
 
     Returns:
        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.
+  def update_training_weight(
+      self, training_weight: tf.Tensor, tensor: tf.Tensor):
+    """Update a training weight to a given tensor value.
+
+    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.
 
     Args:
-       original_layer: tf.keras.layers.Layer representing a layer from the
-       original model.
+      training_weight: tf.Tensor representing a training weight.
+      tensor: tf.Tensor representing a value to be assigned to the training
+        weight.
+    """
+
+  def compress_training_weights(
+      self, *training_weights: tf.Tensor) -> List[tf.Tensor]:
+    """Define the operations to compress a single weight’s training form.
+
+    'compress_training_weights' can refer to making the weight more amenable to
+    compression or actually compress the weight.
+
+    The default is an identity.
+
+    Args:
+      *training_weights: tf.Tensors representing all variables used during
+        training, for a single compressible weight, in the order returned in
+        `init_training_weights`.
 
     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.)
+      List of tf.Tensors to set to compressed or more compressible form.
+    """
+    return list(training_weights)
+
+  @abc.abstractmethod
+  def decompress_weights(
+      self, *compressed_weights: tf.Tensor) -> tf.Tensor:
+    """Define the operations to decompress a single weight’s compressed form.
+
+    The default is an identity.
+
+    Args:
+       *compressed_weights: tf.Tensors representing a single weight’s compressed
+         form, coming from what’s returned in `compress`.
+
+    Returns:
+      A tf.Tensor representing the decompressed `compressed_weights`.
     """
-    del original_layer
-    return []
 
 
 def create_layer_for_training(
     layer: tf.keras.layers.Layer,
-    algorithm: WeightCompressionAlgorithm) -> tf.keras.layers.Layer:
+    algorithm: WeightCompressor) -> tf.keras.layers.Layer:
   return optimize.create_layer_for_training(layer, algorithm)
 
 
 def create_layer_for_inference(
     layer_for_training: tf.keras.layers.Layer,
-    algorithm: WeightCompressionAlgorithm) -> tf.keras.layers.Layer:
+    algorithm: WeightCompressor) -> tf.keras.layers.Layer:
   return optimize.create_layer_for_inference(layer_for_training, algorithm)

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

@@ -23,58 +23,64 @@
 # 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):
+class BiasOnly(algorithm.WeightCompressor):
   """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]:
+  def init_training_weights(
+      self, pretrained_weight: tf.Tensor):
     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))
-    ]
+    self.add_training_weight(
+        name='bias_mean',
+        shape=bias_mean.shape,
+        dtype=bias_mean.dtype,
+        initializer=tf.keras.initializers.Constant(bias_mean))
+    self.add_training_weight(
+        name='bias_shape',
+        shape=bias_shape.shape,
+        dtype=bias_shape.dtype,
+        initializer=tf.keras.initializers.Constant(bias_shape))
 
-  def decompress(
+  def decompress_weights(
       self, bias_mean: tf.Tensor, bias_shape: tf.Tensor) -> tf.Tensor:
     return tf.broadcast_to(bias_mean, bias_shape)
 
-  def training(
+  def project_training_weights(
       self, bias_mean: tf.Tensor, bias_shape: tf.Tensor) -> tf.Tensor:
-    return self.decompress(bias_mean, bias_shape)
+    return self.decompress_weights(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 [original_layer.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:
+  def compress_model(self, to_optimize: tf.keras.Model) -> tf.keras.Model:
+    """Model developer API for optimizing a model."""
+    # pylint: disable=protected-access
+    if not isinstance(to_optimize, tf.keras.Sequential) \
+        and not to_optimize._is_graph_network:
       raise ValueError(
-          'Applying BiasOnly currently requires passing in a built model')
+          '`compress_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 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 algorithm.create_layer_for_training(layer, algorithm=self)
 
-  return tf.keras.models.clone_model(
-      to_optimize, clone_function=_optimize_layer)
+    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

@@ -136,15 +136,15 @@ class FunctionalTest(tf.test.TestCase):
 
   def testBiasOnly_ReducesParamaters(self):
     model = _build_model()
-    compressed_model = bias_only.optimize(model)
+    compressed_model = bias_only.BiasOnly().compress_model(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)
+    compressed_model = bias_only.BiasOnly().compress_model(model)
 
     _train_model(compressed_model)
 
@@ -162,7 +162,7 @@ def testBiasOnly_HasReasonableAccuracy_TF(self):
   def testBiasOnly_HasReasonableAccuracy_TFLite(self):
     model = _build_model()
 
-    compressed_model = bias_only.optimize(model)
+    compressed_model = bias_only.BiasOnly().compress_model(model)
 
     _train_model(compressed_model)
 
@@ -180,7 +180,7 @@ def testBiasOnly_BreaksDownLayerWeights(self):
     first_conv_layer = model.layers[2]
     self.assertLen(first_conv_layer.weights, 2)
 
-    compressed_model = bias_only.optimize(model)
+    compressed_model = bias_only.BiasOnly().compress_model(model)
 
     first_conv_layer = compressed_model.layers[2]
 
@@ -194,20 +194,23 @@ def testBiasOnly_PreservesPretrainedWeights(self):
 
     dense_layer_weights = model.layers[1].get_weights()
 
-    compressed_model = bias_only.optimize(model)
+    algorithm = bias_only.BiasOnly()
+    compressed_model = algorithm.compress_model(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)
+    algorithm.weight_reprs = []
+    algorithm.init_training_weights(dense_layer_weights[1])
+    w1_repr, w2_repr = algorithm.weight_reprs
+
+    w1 = w1_repr.kwargs['initializer'](
+        shape=None, dtype=w1_repr.kwargs['dtype'])
+    w2 = w2_repr.kwargs['initializer'](
+        shape=None, dtype=w2_repr.kwargs['dtype'])
 
     assert (w1 == dense_layer_compressed_weights[0]).numpy().all()
     assert (w2 == dense_layer_compressed_weights[1]).numpy().all()