Update all keras references to use the keras v2 version instance

PiperOrigin-RevId: 603685481

ci/kokoro/gcp_ubuntu/Dockerfile

@@ -26,10 +26,10 @@
 # run CI against.
 
 # Latest Ubuntu LTS (Focal), at the moment.
-FROM ubuntu:20.04
+FROM ubuntu:22.04
 
-ARG BAZEL_VERSION=4.2.2
-ARG TENSORFLOW_VERSION=2.12.0
+ARG BAZEL_VERSION=7.0.2
+ARG TENSORFLOW_VERSION=2.15.0
 
 
 RUN apt-get update -y

ci/kokoro/run_bazel_unittests.sh

@@ -33,9 +33,7 @@ set -o pipefail  # Treat the failure of a command in a pipeline as error.
 # set -x
 
 pip install --requirement "requirements.txt"
-# TODO(b/232345872): Not in list of requirements, but needed for EPR test.
-# The EPR test relies on a feature (PowerLawEntropyModel) introduced in 2.10.0.
-pip install tensorflow-compression~=2.11.0
+pip install tensorflow-compression>=2.11.0
 
 # Run the tests.
 # Some tests requiring more RAM than the CI machine provides are disabled.

tensorflow_model_optimization/g3doc/_index.yaml

@@ -56,8 +56,9 @@ landing_page:
         <pre class = "prettyprint">
         import tensorflow as tf
         import tensorflow_model_optimization as tfmot
+        import tf_keras as keras
 
-        model = tf.keras.Sequential([...])
+        model = keras.Sequential([...])
 
         pruning_schedule = tfmot.sparsity.keras.PolynomialDecay(
                               initial_sparsity=0.0, final_sparsity=0.5,

tensorflow_model_optimization/g3doc/guide/clustering/clustering_comprehensive_guide.ipynb

@@ -102,6 +102,7 @@
         "! pip install -q tensorflow-model-optimization\n",
         "\n",
         "import tensorflow as tf\n",
+        "import tf_keras as keras\n",
         "import numpy as np\n",
         "import tempfile\n",
         "import os\n",
@@ -110,18 +111,18 @@
         "input_dim = 20\n",
         "output_dim = 20\n",
         "x_train = np.random.randn(1, input_dim).astype(np.float32)\n",
-        "y_train = tf.keras.utils.to_categorical(np.random.randn(1), num_classes=output_dim)\n",
+        "y_train = keras.utils.to_categorical(np.random.randn(1), num_classes=output_dim)\n",
         "\n",
         "def setup_model():\n",
-        "  model = tf.keras.Sequential([\n",
-        "      tf.keras.layers.Dense(input_dim, input_shape=[input_dim]),\n",
-        "      tf.keras.layers.Flatten()\n",
+        "  model = keras.Sequential([\n",
+        "      keras.layers.Dense(input_dim, input_shape=[input_dim]),\n",
+        "      keras.layers.Flatten()\n",
         "  ])\n",
         "  return model\n",
         "\n",
         "def train_model(model):\n",
         "  model.compile(\n",
-        "      loss=tf.keras.losses.categorical_crossentropy,\n",
+        "      loss=keras.losses.categorical_crossentropy,\n",
         "      optimizer='adam',\n",
         "      metrics=['accuracy']\n",
         "  )\n",
@@ -243,7 +244,7 @@
         "**Tips** for better model accuracy:\n",
         "\n",
         "* You must pass a pre-trained model with acceptable accuracy to this API. Training models from scratch with clustering results in subpar accuracy.\n",
-        "* Cluster later layers with more redundant parameters (e.g. `tf.keras.layers.Dense`, `tf.keras.layers.Conv2D`), as opposed to the early layers.\n",
+        "* Cluster later layers with more redundant parameters (e.g. `keras.layers.Dense`, `keras.layers.Conv2D`), as opposed to the early layers.\n",
         "* Freeze early layers prior to the clustered layers during fine-tuning. Treat the number of frozen layers as a hyperparameter. Empirically, freezing most early layers is ideal for the current clustering API.\n",
         "* Avoid clustering critical layers (e.g. attention mechanism).\n",
         "\n",
@@ -265,13 +266,13 @@
         "# Helper function uses `cluster_weights` to make only \n",
         "# the Dense layers train with clustering\n",
         "def apply_clustering_to_dense(layer):\n",
-        "  if isinstance(layer, tf.keras.layers.Dense):\n",
+        "  if isinstance(layer, keras.layers.Dense):\n",
         "    return cluster_weights(layer, **clustering_params)\n",
         "  return layer\n",
         "\n",
-        "# Use `tf.keras.models.clone_model` to apply `apply_clustering_to_dense` \n",
+        "# Use `keras.models.clone_model` to apply `apply_clustering_to_dense` \n",
         "# to the layers of the model.\n",
-        "clustered_model = tf.keras.models.clone_model(\n",
+        "clustered_model = keras.models.clone_model(\n",
         "    base_model,\n",
         "    clone_function=apply_clustering_to_dense,\n",
         ")\n",
@@ -326,17 +327,17 @@
       },
       "outputs": [],
       "source": [
-        "class MyDenseLayer(tf.keras.layers.Dense, tfmot.clustering.keras.ClusterableLayer):\n",
+        "class MyDenseLayer(keras.layers.Dense, tfmot.clustering.keras.ClusterableLayer):\n",
         "\n",
         "  def get_clusterable_weights(self):\n",
         "   # Cluster kernel and bias. This is just an example, clustering\n",
         "   # bias usually hurts model accuracy.\n",
         "   return [('kernel', self.kernel), ('bias', self.bias)]\n",
         "\n",
         "# Use `cluster_weights` to make the `MyDenseLayer` layer train with clustering as usual.\n",
-        "model_for_clustering = tf.keras.Sequential([\n",
+        "model_for_clustering = keras.Sequential([\n",
         "  tfmot.clustering.keras.cluster_weights(MyDenseLayer(20, input_shape=[input_dim]), **clustering_params),\n",
-        "  tf.keras.layers.Flatten()\n",
+        "  keras.layers.Flatten()\n",
         "])\n",
         "\n",
         "model_for_clustering.summary()"
@@ -348,7 +349,7 @@
         "id": "SYlWPXEWmxTs"
       },
       "source": [
-        "You may also use `tfmot.clustering.keras.ClusterableLayer` to cluster a keras custom layer. To do this, you extend `tf.keras.Layer` as usual and implement the `__init__`, `call`, and `build` functions, but you also need to extend the `clusterable_layer.ClusterableLayer` class and implement:\n",
+        "You may also use `tfmot.clustering.keras.ClusterableLayer` to cluster a keras custom layer. To do this, you extend `keras.Layer` as usual and implement the `__init__`, `call`, and `build` functions, but you also need to extend the `clusterable_layer.ClusterableLayer` class and implement:\n",
         "1. `get_clusterable_weights`, where you specify the weight kernel to be clustered, as shown above.\n",
         "2. `get_clusterable_algorithm`, where you specify the clustering algorithm for the weight tensor. This is because you need to specify how the custom layer weights are shaped for clustering. The returned clustering algorithm class should be derived from the `clustering_algorithm.ClusteringAlgorithm` class and the function `get_pulling_indices` should be overwritten. An example of this function, which supports weights of ranks 1D, 2D, and 3D, can be found [here]( https://github.com/tensorflow/model-optimization/blob/18e87d262e536c9a742aef700880e71b47a7f768/tensorflow_model_optimization/python/core/clustering/keras/clustering_algorithm.py#L62).\n",
         "\n",
@@ -392,7 +393,7 @@
         "\n",
         "# `cluster_scope` is needed for deserializing HDF5 models.\n",
         "with tfmot.clustering.keras.cluster_scope():\n",
-        "  loaded_model = tf.keras.models.load_model(keras_model_file)\n",
+        "  loaded_model = keras.models.load_model(keras_model_file)\n",
         "\n",
         "loaded_model.summary()"
       ]
@@ -460,7 +461,7 @@
         "clustered_model = cluster_weights(model, **clustering_params)\n",
         "\n",
         "clustered_model.compile(\n",
-        "    loss=tf.keras.losses.categorical_crossentropy,\n",
+        "    loss=keras.losses.categorical_crossentropy,\n",
         "    optimizer='adam',\n",
         "    metrics=['accuracy']\n",
         ")\n",

tensorflow_model_optimization/g3doc/guide/clustering/clustering_example.ipynb

@@ -82,7 +82,7 @@
         "\n",
         "In the tutorial, you will:\n",
         "\n",
-        "1. Train a `tf.keras` model for the MNIST dataset from scratch.\n",
+        "1. Train a `keras` model for the MNIST dataset from scratch.\n",
         "2. Fine-tune the model by applying the weight clustering API and see the accuracy.\n",
         "3. Create a 6x smaller TF and TFLite models from clustering.\n",
         "4. Create a 8x smaller TFLite model from combining weight clustering and post-training quantization.\n",
@@ -120,7 +120,7 @@
       "outputs": [],
       "source": [
         "import tensorflow as tf\n",
-        "from tensorflow import keras\n",
+        "import tf_keras as keras\n",
         "\n",
         "import numpy as np\n",
         "import tempfile\n",
@@ -134,7 +134,7 @@
         "id": "dKzOfl5FSGPL"
       },
       "source": [
-        "## Train a tf.keras model for MNIST without clustering"
+        "## Train a keras model for MNIST without clustering"
       ]
     },
     {
@@ -146,8 +146,7 @@
       "outputs": [],
       "source": [
         "# Load MNIST dataset\n",
-        "mnist = keras.datasets.mnist\n",
-        "(train_images, train_labels), (test_images, test_labels) = mnist.load_data()\n",
+        "(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()\n",
         "\n",
         "# Normalize the input image so that each pixel value is between 0 to 1.\n",
         "train_images = train_images / 255.0\n",
@@ -165,7 +164,7 @@
         "\n",
         "# Train the digit classification model\n",
         "model.compile(optimizer='adam',\n",
-        "              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "              metrics=['accuracy'])\n",
         "\n",
         "model.fit(\n",
@@ -200,7 +199,7 @@
         "\n",
         "_, keras_file = tempfile.mkstemp('.h5')\n",
         "print('Saving model to: ', keras_file)\n",
-        "tf.keras.models.save_model(model, keras_file, include_optimizer=False)"
+        "keras.models.save_model(model, keras_file, include_optimizer=False)"
       ]
     },
     {
@@ -261,10 +260,10 @@
         "clustered_model = cluster_weights(model, **clustering_params)\n",
         "\n",
         "# Use smaller learning rate for fine-tuning clustered model\n",
-        "opt = tf.keras.optimizers.Adam(learning_rate=1e-5)\n",
+        "opt = keras.optimizers.Adam(learning_rate=1e-5)\n",
         "\n",
         "clustered_model.compile(\n",
-        "  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "  loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "  optimizer=opt,\n",
         "  metrics=['accuracy'])\n",
         "\n",
@@ -362,7 +361,7 @@
         "\n",
         "_, clustered_keras_file = tempfile.mkstemp('.h5')\n",
         "print('Saving clustered model to: ', clustered_keras_file)\n",
-        "tf.keras.models.save_model(final_model, clustered_keras_file, \n",
+        "keras.models.save_model(final_model, clustered_keras_file, \n",
         "                           include_optimizer=False)"
       ]
     },

tensorflow_model_optimization/g3doc/guide/clustering/index.md

@@ -21,7 +21,7 @@ Please note that clustering will provide reduced benefits for convolution and de
 
 Users can apply clustering with the following APIs:
 
-*   Model building: `tf.keras` with only Sequential and Functional models
+*   Model building: `keras` with only Sequential and Functional models
 *   TensorFlow versions: TF 1.x for versions 1.14+ and 2.x.
     *   `tf.compat.v1` with a TF 2.X package and `tf.compat.v2` with a TF 1.X
         package are not supported.

tensorflow_model_optimization/g3doc/guide/combine/cqat_example.ipynb

@@ -80,7 +80,7 @@
         "\n",
         "In the tutorial, you will:\n",
         "\n",
-        "1. Train a `tf.keras` model for the MNIST dataset from scratch.\n",
+        "1. Train a `keras` model for the MNIST dataset from scratch.\n",
         "2. Fine-tune the model with clustering and see the accuracy.\n",
         "3. Apply QAT and observe the loss of clusters.\n",
         "4. Apply CQAT and observe that the clustering applied earlier has been preserved.\n",
@@ -119,6 +119,7 @@
       "outputs": [],
       "source": [
         "import tensorflow as tf\n",
+        "import tf_keras as keras\n",
         "\n",
         "import numpy as np\n",
         "import tempfile\n",
@@ -132,7 +133,7 @@
         "id": "dKzOfl5FSGPL"
       },
       "source": [
-        "## Train a tf.keras model for MNIST without clustering"
+        "## Train a keras model for MNIST without clustering"
       ]
     },
     {
@@ -144,26 +145,26 @@
       "outputs": [],
       "source": [
         "# Load MNIST dataset\n",
-        "mnist = tf.keras.datasets.mnist\n",
+        "mnist = keras.datasets.mnist\n",
         "(train_images, train_labels), (test_images, test_labels) = mnist.load_data()\n",
         "\n",
         "# Normalize the input image so that each pixel value is between 0 to 1.\n",
         "train_images = train_images / 255.0\n",
         "test_images  = test_images / 255.0\n",
         "\n",
-        "model = tf.keras.Sequential([\n",
-        "  tf.keras.layers.InputLayer(input_shape=(28, 28)),\n",
-        "  tf.keras.layers.Reshape(target_shape=(28, 28, 1)),\n",
-        "  tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3),\n",
+        "model = keras.Sequential([\n",
+        "  keras.layers.InputLayer(input_shape=(28, 28)),\n",
+        "  keras.layers.Reshape(target_shape=(28, 28, 1)),\n",
+        "  keras.layers.Conv2D(filters=12, kernel_size=(3, 3),\n",
         "                         activation=tf.nn.relu),\n",
-        "  tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
-        "  tf.keras.layers.Flatten(),\n",
-        "  tf.keras.layers.Dense(10)\n",
+        "  keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
+        "  keras.layers.Flatten(),\n",
+        "  keras.layers.Dense(10)\n",
         "])\n",
         "\n",
         "# Train the digit classification model\n",
         "model.compile(optimizer='adam',\n",
-        "              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "              metrics=['accuracy'])\n",
         "\n",
         "model.fit(\n",
@@ -198,7 +199,7 @@
         "\n",
         "_, keras_file = tempfile.mkstemp('.h5')\n",
         "print('Saving model to: ', keras_file)\n",
-        "tf.keras.models.save_model(model, keras_file, include_optimizer=False)"
+        "keras.models.save_model(model, keras_file, include_optimizer=False)"
       ]
     },
     {
@@ -259,10 +260,10 @@
         "clustered_model = cluster_weights(model, **clustering_params)\n",
         "\n",
         "# Use smaller learning rate for fine-tuning\n",
-        "opt = tf.keras.optimizers.Adam(learning_rate=1e-5)\n",
+        "opt = keras.optimizers.Adam(learning_rate=1e-5)\n",
         "\n",
         "clustered_model.compile(\n",
-        "  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "  loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "  optimizer=opt,\n",
         "  metrics=['accuracy'])\n",
         "\n",
@@ -323,7 +324,7 @@
         "def print_model_weight_clusters(model):\n",
         "\n",
         "    for layer in model.layers:\n",
-        "        if isinstance(layer, tf.keras.layers.Wrapper):\n",
+        "        if isinstance(layer, keras.layers.Wrapper):\n",
         "            weights = layer.trainable_weights\n",
         "        else:\n",
         "            weights = layer.weights\n",
@@ -414,7 +415,7 @@
         "qat_model = tfmot.quantization.keras.quantize_model(stripped_clustered_model)\n",
         "\n",
         "qat_model.compile(optimizer='adam',\n",
-        "              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "              metrics=['accuracy'])\n",
         "print('Train qat model:')\n",
         "qat_model.fit(train_images, train_labels, batch_size=128, epochs=1, validation_split=0.1)\n",
@@ -427,7 +428,7 @@
         "              tfmot.experimental.combine.Default8BitClusterPreserveQuantizeScheme())\n",
         "\n",
         "cqat_model.compile(optimizer='adam',\n",
-        "              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "              metrics=['accuracy'])\n",
         "print('Train cqat model:')\n",
         "cqat_model.fit(train_images, train_labels, batch_size=128, epochs=1, validation_split=0.1)"