Merge pull request #952 from wwwind:clustering_conv1d_transpose

PiperOrigin-RevId: 448273933

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/clustering/keras/cluster_integration_test.py

@@ -85,6 +85,11 @@ def dataset_generator2(self):
     for x, y in zip(self.x_train2, self.y_train2):
       yield np.array([x]), np.array([y])
 
+  def _batch(self, dims, batch_size):
+    if dims[0] is None:
+      dims[0] = batch_size
+    return dims
+
   def end_to_end_testing(self, original_model, clusters_check=None):
     """Test End to End clustering."""
 
@@ -225,6 +230,80 @@ def testSparsityIsPreservedDuringTraining(self):
         nr_of_unique_weights_after,
         clustering_params["number_of_clusters"])
 
+  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
+  def testEndToEndSequential(self):
+    """Test End to End clustering - sequential model."""
+    original_model = keras.Sequential([
+        layers.Dense(5, input_shape=(5,)),
+        layers.Dense(5),
+    ])
+
+    def clusters_check(stripped_model):
+      # dense layer
+      weights_as_list = stripped_model.get_weights()[0].reshape(-1,).tolist()
+      unique_weights = set(weights_as_list)
+      self.assertLessEqual(
+          len(unique_weights), self.params["number_of_clusters"])
+
+    self.end_to_end_testing(original_model, clusters_check)
+
+  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
+  def testEndToEndConv1DAndConv1DTranspose(self):
+    """Test End to End clustering - model with Conv1D and Conv1DTranspose."""
+    inp = layers.Input(batch_shape=(1, 16))
+    x = layers.Conv1D(
+        10, 16, 4, padding="valid", use_bias=False)(
+            tf.expand_dims(inp, axis=-1))
+    y = layers.Conv1DTranspose(1, 16, 4, padding="valid", use_bias=False)(x)
+    model = keras.models.Model(inputs=inp, outputs=[y])
+
+    def apply_clustering(layer):
+      if isinstance(layer, keras.layers.Conv1D) or isinstance(
+          layer, keras.layers.Conv1DTranspose):
+        return cluster.cluster_weights(layer, **self.params)
+      return layer
+
+    model_to_cluster = keras.models.clone_model(
+        model,
+        clone_function=apply_clustering,
+    )
+
+    model_to_cluster.compile(
+        loss=keras.losses.categorical_crossentropy,
+        optimizer="adam",
+        metrics=["accuracy"]
+    )
+    model_to_cluster.fit(
+        np.random.randn(*self._batch(model.input.get_shape().as_list(), 16)),
+        np.random.randn(*self._batch(model.output.get_shape().as_list(), 16)),
+        steps_per_epoch=1)
+    clustered_model = cluster.strip_clustering(model_to_cluster)
+
+    def do_checks(layer, layer_name):
+      self.assertEqual(layer.name, layer_name)
+      unique_weights = np.unique(layer.weights[0].numpy().flatten())
+      self.assertLessEqual(
+          len(unique_weights), self.params["number_of_clusters"])
+
+    do_checks(clustered_model.layers[2], "conv1d")
+    do_checks(clustered_model.layers[3], "conv1d_transpose")
+
+  def testStripClusteringSequentialModelWithRegulariser(self):
+    """Verifies that stripping the clustering wrappers from a sequential model produces the expected config."""
+    original_model = keras.Sequential([
+        layers.Dense(5, input_shape=(5,)),
+        layers.Dense(5, kernel_regularizer=tf.keras.regularizers.L1(0.01)),
+    ])
+
+    def clusters_check(stripped_model):
+      # dense layer
+      weights_as_list = stripped_model.get_weights()[0].reshape(-1,).tolist()
+      unique_weights = set(weights_as_list)
+      self.assertLessEqual(
+          len(unique_weights), self.params["number_of_clusters"])
+
+    self.end_to_end_testing(original_model, clusters_check)
+
   @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
   def testEndToEndFunctional(self):
     """Test End to End clustering - functional model."""

tensorflow_model_optimization/python/core/clustering/keras/cluster_test.py

@@ -127,7 +127,11 @@ def setUp(self):
     self.keras_depthwiseconv2d_layer = layers.DepthwiseConv2D((3, 3), (1, 1))
     self.keras_dense_layer = layers.Dense(10)
     self.keras_conv1d_layer = layers.Conv1D(filters=3, kernel_size=(5))
+    self.keras_conv1d_tr_layer = layers.Conv1DTranspose(
+        filters=3, kernel_size=(5))
     self.keras_conv2d_layer = layers.Conv2D(filters=3, kernel_size=(4, 5))
+    self.keras_conv2d_tr_layer = layers.Conv2DTranspose(
+        filters=3, kernel_size=(4, 5))
     self.keras_conv3d_layer = layers.Conv3D(filters=2, kernel_size=(3, 4, 5))
     self.keras_custom_layer = KerasCustomLayer()
     self.clusterable_layer = MyClusterableLayer(10)
@@ -223,6 +227,19 @@ def testConv1DLayer(self):
     self.assertEqual([5, 1, 3],
                      wrapped_layer.layer.get_clusterable_weights()[0][1].shape)
 
+  @keras_parameterized.run_all_keras_modes
+  def testConv1DTransposeLayer(self):
+    """Verifies that we can cluster a Conv1DTranspose layer."""
+    input_shape = (4, 28, 1)
+    wrapped_layer = self._build_clustered_layer_model(
+        self.keras_conv1d_tr_layer,
+        input_shape=input_shape)
+
+    self._validate_clustered_layer(self.keras_conv1d_tr_layer,
+                                   wrapped_layer)
+    self.assertEqual([5, 3, 1],
+                     wrapped_layer.layer.get_clusterable_weights()[0][1].shape)
+
   @keras_parameterized.run_all_keras_modes
   def testConv2DLayer(self):
     """Verifies that we can cluster a Conv2D layer."""
@@ -236,6 +253,19 @@ def testConv2DLayer(self):
     self.assertEqual([4, 5, 1, 3],
                      wrapped_layer.layer.get_clusterable_weights()[0][1].shape)
 
+  @keras_parameterized.run_all_keras_modes
+  def testConv2DTransposeLayer(self):
+    """Verifies that we can cluster a Conv2DTranspose layer."""
+    input_shape = (4, 28, 28, 1)
+    wrapped_layer = self._build_clustered_layer_model(
+        self.keras_conv2d_tr_layer,
+        input_shape=input_shape)
+
+    self._validate_clustered_layer(self.keras_conv2d_tr_layer,
+                                   wrapped_layer)
+    self.assertEqual([4, 5, 3, 1],
+                     wrapped_layer.layer.get_clusterable_weights()[0][1].shape)
+
   @keras_parameterized.run_all_keras_modes
   def testConv3DLayer(self):
     """Verifies that we can cluster a Conv3D layer."""

tensorflow_model_optimization/python/core/clustering/keras/clustering_registry.py

@@ -62,6 +62,7 @@ class ClusteringRegistry(object):
   # allows the wrapper to access and modify the weights.
   _LAYERS_WEIGHTS_MAP = {
       layers.Conv1D: ['kernel'],
+      layers.Conv1DTranspose: ['kernel'],
       layers.Conv2D: ['kernel'],
       layers.Conv2DTranspose: ['kernel'],
       layers.Conv3D: ['kernel'],