Remove explicit line "experimental_new_converter = True" the converter launched and is now the default.

PiperOrigin-RevId: 318582066

Author: karimnosseir

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

@@ -14,123 +14,119 @@
 # ==============================================================================
 """End-to-end tests for keras clustering API."""
 
-import numpy as np
-import tensorflow as tf
+import os
 import tempfile
+
 from absl.testing import parameterized
-from tensorflow.python.keras import keras_parameterized
+import numpy as np
+import tensorflow as tf
 
+from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.clustering.keras import cluster
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 from tensorflow_model_optimization.python.core.keras import compat
-import os
 
 keras = tf.keras
 layers = keras.layers
 test = tf.test
 
 CentroidInitialization = cluster_config.CentroidInitialization
 
+
 class ClusterIntegrationTest(test.TestCase, parameterized.TestCase):
+  """Integration tests for clustering."""
+
+  def setUp(self):
+    self.params = {
+        "number_of_clusters": 8,
+        "cluster_centroids_init": CentroidInitialization.LINEAR,
+    }
+
+    self.x_train = np.array(
+        [[0.0, 1.0], [2.0, 0.0], [0.0, 3.0], [4.0, 1.0], [5.0, 1.0]],
+        dtype="float32",
+    )
+
+    self.y_train = np.array(
+        [[0.0, 1.0], [1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]],
+        dtype="float32",
+    )
+
+  def dataset_generator(self):
+    for x, y in zip(self.x_train, self.y_train):
+      yield np.array([x]), np.array([y])
+
+  @staticmethod
+  def _verify_tflite(tflite_file, x_test):
+    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 = x_test[0]
+    x = x.reshape((1,) + x.shape)
+    interpreter.set_tensor(input_index, x)
+    interpreter.invoke()
+    interpreter.get_tensor(output_index)
+
+  @keras_parameterized.run_all_keras_modes
+  def testValuesRemainClusteredAfterTraining(self):
+    """Verifies that training a clustered model does not destroy the clusters."""
+    original_model = keras.Sequential([
+        layers.Dense(2, input_shape=(2,)),
+        layers.Dense(2),
+    ])
+
+    clustered_model = cluster.cluster_weights(original_model, **self.params)
+
+    clustered_model.compile(
+        loss=keras.losses.categorical_crossentropy,
+        optimizer="adam",
+        metrics=["accuracy"],
+    )
+
+    clustered_model.fit(x=self.dataset_generator(), steps_per_epoch=1)
+    stripped_model = cluster.strip_clustering(clustered_model)
+    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"])
+
+  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
+  def testEndToEnd(self):
+    """Test End to End clustering."""
+    original_model = keras.Sequential([
+        layers.Dense(2, input_shape=(2,)),
+        layers.Dense(2),
+    ])
+
+    clustered_model = cluster.cluster_weights(original_model, **self.params)
+
+    clustered_model.compile(
+        loss=keras.losses.categorical_crossentropy,
+        optimizer="adam",
+        metrics=["accuracy"],
+    )
+
+    clustered_model.fit(x=self.dataset_generator(), steps_per_epoch=1)
+    stripped_model = cluster.strip_clustering(clustered_model)
+
+    _, tflite_file = tempfile.mkstemp(".tflite")
+    _, keras_file = tempfile.mkstemp(".h5")
+
+    if not compat.is_v1_apis():
+      converter = tf.lite.TFLiteConverter.from_keras_model(stripped_model)
+    else:
+      tf.keras.models.save_model(stripped_model, keras_file)
+      converter = tf.lite.TFLiteConverter.from_keras_model_file(keras_file)
+
+    tflite_model = converter.convert()
+    with open(tflite_file, "wb") as f:
+      f.write(tflite_model)
+
+    self._verify_tflite(tflite_file, self.x_train)
+
+    os.remove(keras_file)
+    os.remove(tflite_file)
 
-    """
-    Integration tests for clustering.
-    """
-    def setUp(self):
-        self.params = {
-            "number_of_clusters": 8,
-            "cluster_centroids_init": CentroidInitialization.LINEAR,
-        }
-
-        self.x_train = np.array(
-            [[0.0, 1.0], [2.0, 0.0], [0.0, 3.0], [4.0, 1.0], [5.0, 1.0]],
-            dtype="float32",
-        )
-
-        self.y_train = np.array(
-            [[0.0, 1.0], [1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]],
-            dtype="float32",
-        )
-
-    def dataset_generator(self):
-        for x, y in zip(self.x_train, self.y_train):
-            yield np.array([x]), np.array([y])
-
-    @staticmethod
-    def _verify_tflite(tflite_file, x_test):
-        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 = x_test[0]
-        x = x.reshape((1,) + x.shape)
-        interpreter.set_tensor(input_index, x)
-        interpreter.invoke()
-        interpreter.get_tensor(output_index)
-
-    @keras_parameterized.run_all_keras_modes
-    def testValuesRemainClusteredAfterTraining(self):
-
-        """
-        Verifies that training a clustered model does not destroy the clusters.
-        """
-        original_model = keras.Sequential(
-            [layers.Dense(2, input_shape=(2,)), layers.Dense(2),]
-        )
-
-        clustered_model = cluster.cluster_weights(original_model, **self.params)
-
-        clustered_model.compile(
-            loss=keras.losses.categorical_crossentropy,
-            optimizer="adam",
-            metrics=["accuracy"],
-        )
-
-        clustered_model.fit(x=self.dataset_generator(), steps_per_epoch=1)
-        stripped_model = cluster.strip_clustering(clustered_model)
-        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"])
-
-    @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
-    def testEndToEnd(self):
-
-        """
-        Test End to End clustering.
-        """
-        original_model = keras.Sequential(
-            [layers.Dense(2, input_shape=(2,)), layers.Dense(2),]
-        )
-
-        clustered_model = cluster.cluster_weights(original_model, **self.params)
-
-        clustered_model.compile(
-            loss=keras.losses.categorical_crossentropy,
-            optimizer="adam",
-            metrics=["accuracy"],
-        )
-
-        clustered_model.fit(x=self.dataset_generator(), steps_per_epoch=1)
-        stripped_model = cluster.strip_clustering(clustered_model)
-
-        _, tflite_file = tempfile.mkstemp(".tflite")
-        _, keras_file = tempfile.mkstemp(".h5")
-
-        if not compat.is_v1_apis():
-            converter = tf.lite.TFLiteConverter.from_keras_model(stripped_model)
-        else:
-            tf.keras.models.save_model(stripped_model, keras_file)
-            converter = tf.lite.TFLiteConverter.from_keras_model_file(keras_file)
-
-        converter.experimental_new_converter = True
-        tflite_model = converter.convert()
-        with open(tflite_file, "wb") as f:
-            f.write(tflite_model)
-
-        self._verify_tflite(tflite_file, self.x_train)
-
-        os.remove(keras_file)
-        os.remove(tflite_file)
 
 if __name__ == "__main__":
-    test.main()
+  test.main()

tensorflow_model_optimization/python/examples/sparsity/keras/mnist/mnist_e2e.py

@@ -112,7 +112,6 @@ def main(unused_argv):
   model = train(model, x_train, y_train, x_test, y_test)
 
   converter = tf.lite.TFLiteConverter.from_keras_model(model)
-  converter.experimental_new_converter = True
 
   # Get a dense model as baseline
   tflite_model_dense = converter.convert()
@@ -153,7 +152,6 @@ def main(unused_argv):
   model = train(model, x_train, y_train, x_test, y_test)
 
   converter = tf.lite.TFLiteConverter.from_keras_model(model)
-  converter.experimental_new_converter = True
   converter._experimental_sparsify_model = True
 
   tflite_model = converter.convert()