Skip to content

Commit b0eea48

Browse files
abatterytensorflower-gardener
abattery
authored and
tensorflower-gardener
committed
Update all keras references to use the keras v2 version instance
PiperOrigin-RevId: 603753268
1 parent 9533188 commit b0eea48

File tree

143 files changed

+1958
-1445
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

143 files changed

+1958
-1445
lines changed

‎ci/kokoro/gcp_ubuntu/Dockerfile‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,10 @@
2626
# run CI against.
2727

2828
# Latest Ubuntu LTS (Focal), at the moment.
29-
FROM ubuntu:20.04
29+
FROM ubuntu:22.04
3030

31-
ARG BAZEL_VERSION=4.2.2
32-
ARG TENSORFLOW_VERSION=2.12.0
31+
ARG BAZEL_VERSION=7.0.2
32+
ARG TENSORFLOW_VERSION=2.15.0
3333

3434

3535
RUN apt-get update -y

‎ci/kokoro/run_bazel_unittests.sh‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,7 @@ set -o pipefail # Treat the failure of a command in a pipeline as error.
3333
# set -x
3434

3535
pip install --requirement "requirements.txt"
36-
# TODO(b/232345872): Not in list of requirements, but needed for EPR test.
37-
# The EPR test relies on a feature (PowerLawEntropyModel) introduced in 2.10.0.
38-
pip install tensorflow-compression~=2.11.0
36+
pip install tensorflow-compression>=2.11.0
3937

4038
# Run the tests.
4139
# Some tests requiring more RAM than the CI machine provides are disabled.

‎tensorflow_model_optimization/g3doc/_index.yaml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,8 +56,9 @@ landing_page:
5656
<pre class = "prettyprint">
5757
import tensorflow as tf
5858
import tensorflow_model_optimization as tfmot
59+
import tf_keras as keras
5960
60-
model = tf.keras.Sequential([...])
61+
model = keras.Sequential([...])
6162
6263
pruning_schedule = tfmot.sparsity.keras.PolynomialDecay(
6364
initial_sparsity=0.0, final_sparsity=0.5,

‎tensorflow_model_optimization/g3doc/guide/clustering/clustering_comprehensive_guide.ipynb‎

Lines changed: 16 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@
102102
"! pip install -q tensorflow-model-optimization\n",
103103
"\n",
104104
"import tensorflow as tf\n",
105+
"import tf_keras as keras\n",
105106
"import numpy as np\n",
106107
"import tempfile\n",
107108
"import os\n",
@@ -110,18 +111,18 @@
110111
"input_dim = 20\n",
111112
"output_dim = 20\n",
112113
"x_train = np.random.randn(1, input_dim).astype(np.float32)\n",
113-
"y_train = tf.keras.utils.to_categorical(np.random.randn(1), num_classes=output_dim)\n",
114+
"y_train = keras.utils.to_categorical(np.random.randn(1), num_classes=output_dim)\n",
114115
"\n",
115116
"def setup_model():\n",
116-
" model = tf.keras.Sequential([\n",
117-
" tf.keras.layers.Dense(input_dim, input_shape=[input_dim]),\n",
118-
" tf.keras.layers.Flatten()\n",
117+
" model = keras.Sequential([\n",
118+
" keras.layers.Dense(input_dim, input_shape=[input_dim]),\n",
119+
" keras.layers.Flatten()\n",
119120
" ])\n",
120121
" return model\n",
121122
"\n",
122123
"def train_model(model):\n",
123124
" model.compile(\n",
124-
" loss=tf.keras.losses.categorical_crossentropy,\n",
125+
" loss=keras.losses.categorical_crossentropy,\n",
125126
" optimizer='adam',\n",
126127
" metrics=['accuracy']\n",
127128
" )\n",
@@ -243,7 +244,7 @@
243244
"**Tips** for better model accuracy:\n",
244245
"\n",
245246
"* You must pass a pre-trained model with acceptable accuracy to this API. Training models from scratch with clustering results in subpar accuracy.\n",
246-
"* Cluster later layers with more redundant parameters (e.g. `tf.keras.layers.Dense`, `tf.keras.layers.Conv2D`), as opposed to the early layers.\n",
247+
"* Cluster later layers with more redundant parameters (e.g. `keras.layers.Dense`, `keras.layers.Conv2D`), as opposed to the early layers.\n",
247248
"* 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",
248249
"* Avoid clustering critical layers (e.g. attention mechanism).\n",
249250
"\n",
@@ -265,13 +266,13 @@
265266
"# Helper function uses `cluster_weights` to make only \n",
266267
"# the Dense layers train with clustering\n",
267268
"def apply_clustering_to_dense(layer):\n",
268-
" if isinstance(layer, tf.keras.layers.Dense):\n",
269+
" if isinstance(layer, keras.layers.Dense):\n",
269270
" return cluster_weights(layer, **clustering_params)\n",
270271
" return layer\n",
271272
"\n",
272-
"# Use `tf.keras.models.clone_model` to apply `apply_clustering_to_dense` \n",
273+
"# Use `keras.models.clone_model` to apply `apply_clustering_to_dense` \n",
273274
"# to the layers of the model.\n",
274-
"clustered_model = tf.keras.models.clone_model(\n",
275+
"clustered_model = keras.models.clone_model(\n",
275276
" base_model,\n",
276277
" clone_function=apply_clustering_to_dense,\n",
277278
")\n",
@@ -326,17 +327,17 @@
326327
},
327328
"outputs": [],
328329
"source": [
329-
"class MyDenseLayer(tf.keras.layers.Dense, tfmot.clustering.keras.ClusterableLayer):\n",
330+
"class MyDenseLayer(keras.layers.Dense, tfmot.clustering.keras.ClusterableLayer):\n",
330331
"\n",
331332
" def get_clusterable_weights(self):\n",
332333
" # Cluster kernel and bias. This is just an example, clustering\n",
333334
" # bias usually hurts model accuracy.\n",
334335
" return [('kernel', self.kernel), ('bias', self.bias)]\n",
335336
"\n",
336337
"# Use `cluster_weights` to make the `MyDenseLayer` layer train with clustering as usual.\n",
337-
"model_for_clustering = tf.keras.Sequential([\n",
338+
"model_for_clustering = keras.Sequential([\n",
338339
" tfmot.clustering.keras.cluster_weights(MyDenseLayer(20, input_shape=[input_dim]), **clustering_params),\n",
339-
" tf.keras.layers.Flatten()\n",
340+
" keras.layers.Flatten()\n",
340341
"])\n",
341342
"\n",
342343
"model_for_clustering.summary()"
@@ -348,7 +349,7 @@
348349
"id": "SYlWPXEWmxTs"
349350
},
350351
"source": [
351-
"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",
352+
"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",
352353
"1. `get_clusterable_weights`, where you specify the weight kernel to be clustered, as shown above.\n",
353354
"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",
354355
"\n",
@@ -392,7 +393,7 @@
392393
"\n",
393394
"# `cluster_scope` is needed for deserializing HDF5 models.\n",
394395
"with tfmot.clustering.keras.cluster_scope():\n",
395-
" loaded_model = tf.keras.models.load_model(keras_model_file)\n",
396+
" loaded_model = keras.models.load_model(keras_model_file)\n",
396397
"\n",
397398
"loaded_model.summary()"
398399
]
@@ -460,7 +461,7 @@
460461
"clustered_model = cluster_weights(model, **clustering_params)\n",
461462
"\n",
462463
"clustered_model.compile(\n",
463-
" loss=tf.keras.losses.categorical_crossentropy,\n",
464+
" loss=keras.losses.categorical_crossentropy,\n",
464465
" optimizer='adam',\n",
465466
" metrics=['accuracy']\n",
466467
")\n",

‎tensorflow_model_optimization/g3doc/guide/clustering/clustering_example.ipynb‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@
8282
"\n",
8383
"In the tutorial, you will:\n",
8484
"\n",
85-
"1. Train a `tf.keras` model for the MNIST dataset from scratch.\n",
85+
"1. Train a `keras` model for the MNIST dataset from scratch.\n",
8686
"2. Fine-tune the model by applying the weight clustering API and see the accuracy.\n",
8787
"3. Create a 6x smaller TF and TFLite models from clustering.\n",
8888
"4. Create a 8x smaller TFLite model from combining weight clustering and post-training quantization.\n",
@@ -120,7 +120,7 @@
120120
"outputs": [],
121121
"source": [
122122
"import tensorflow as tf\n",
123-
"from tensorflow import keras\n",
123+
"import tf_keras as keras\n",
124124
"\n",
125125
"import numpy as np\n",
126126
"import tempfile\n",
@@ -134,7 +134,7 @@
134134
"id": "dKzOfl5FSGPL"
135135
},
136136
"source": [
137-
"## Train a tf.keras model for MNIST without clustering"
137+
"## Train a keras model for MNIST without clustering"
138138
]
139139
},
140140
{
@@ -146,8 +146,7 @@
146146
"outputs": [],
147147
"source": [
148148
"# Load MNIST dataset\n",
149-
"mnist = keras.datasets.mnist\n",
150-
"(train_images, train_labels), (test_images, test_labels) = mnist.load_data()\n",
149+
"(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()\n",
151150
"\n",
152151
"# Normalize the input image so that each pixel value is between 0 to 1.\n",
153152
"train_images = train_images / 255.0\n",
@@ -165,7 +164,7 @@
165164
"\n",
166165
"# Train the digit classification model\n",
167166
"model.compile(optimizer='adam',\n",
168-
" loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
167+
" loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
169168
" metrics=['accuracy'])\n",
170169
"\n",
171170
"model.fit(\n",
@@ -200,7 +199,7 @@
200199
"\n",
201200
"_, keras_file = tempfile.mkstemp('.h5')\n",
202201
"print('Saving model to: ', keras_file)\n",
203-
"tf.keras.models.save_model(model, keras_file, include_optimizer=False)"
202+
"keras.models.save_model(model, keras_file, include_optimizer=False)"
204203
]
205204
},
206205
{
@@ -261,10 +260,10 @@
261260
"clustered_model = cluster_weights(model, **clustering_params)\n",
262261
"\n",
263262
"# Use smaller learning rate for fine-tuning clustered model\n",
264-
"opt = tf.keras.optimizers.Adam(learning_rate=1e-5)\n",
263+
"opt = keras.optimizers.Adam(learning_rate=1e-5)\n",
265264
"\n",
266265
"clustered_model.compile(\n",
267-
" loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
266+
" loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
268267
" optimizer=opt,\n",
269268
" metrics=['accuracy'])\n",
270269
"\n",
@@ -362,7 +361,7 @@
362361
"\n",
363362
"_, clustered_keras_file = tempfile.mkstemp('.h5')\n",
364363
"print('Saving clustered model to: ', clustered_keras_file)\n",
365-
"tf.keras.models.save_model(final_model, clustered_keras_file, \n",
364+
"keras.models.save_model(final_model, clustered_keras_file, \n",
366365
" include_optimizer=False)"
367366
]
368367
},

‎tensorflow_model_optimization/g3doc/guide/clustering/index.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ Please note that clustering will provide reduced benefits for convolution and de
2121

2222
Users can apply clustering with the following APIs:
2323

24-
* Model building: `tf.keras` with only Sequential and Functional models
24+
* Model building: `keras` with only Sequential and Functional models
2525
* TensorFlow versions: TF 1.x for versions 1.14+ and 2.x.
2626
* `tf.compat.v1` with a TF 2.X package and `tf.compat.v2` with a TF 1.X
2727
package are not supported.

‎tensorflow_model_optimization/g3doc/guide/combine/cqat_example.ipynb‎

Lines changed: 18 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@
8080
"\n",
8181
"In the tutorial, you will:\n",
8282
"\n",
83-
"1. Train a `tf.keras` model for the MNIST dataset from scratch.\n",
83+
"1. Train a `keras` model for the MNIST dataset from scratch.\n",
8484
"2. Fine-tune the model with clustering and see the accuracy.\n",
8585
"3. Apply QAT and observe the loss of clusters.\n",
8686
"4. Apply CQAT and observe that the clustering applied earlier has been preserved.\n",
@@ -119,6 +119,7 @@
119119
"outputs": [],
120120
"source": [
121121
"import tensorflow as tf\n",
122+
"import tf_keras as keras\n",
122123
"\n",
123124
"import numpy as np\n",
124125
"import tempfile\n",
@@ -132,7 +133,7 @@
132133
"id": "dKzOfl5FSGPL"
133134
},
134135
"source": [
135-
"## Train a tf.keras model for MNIST without clustering"
136+
"## Train a keras model for MNIST without clustering"
136137
]
137138
},
138139
{
@@ -144,26 +145,26 @@
144145
"outputs": [],
145146
"source": [
146147
"# Load MNIST dataset\n",
147-
"mnist = tf.keras.datasets.mnist\n",
148+
"mnist = keras.datasets.mnist\n",
148149
"(train_images, train_labels), (test_images, test_labels) = mnist.load_data()\n",
149150
"\n",
150151
"# Normalize the input image so that each pixel value is between 0 to 1.\n",
151152
"train_images = train_images / 255.0\n",
152153
"test_images = test_images / 255.0\n",
153154
"\n",
154-
"model = tf.keras.Sequential([\n",
155-
" tf.keras.layers.InputLayer(input_shape=(28, 28)),\n",
156-
" tf.keras.layers.Reshape(target_shape=(28, 28, 1)),\n",
157-
" tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3),\n",
155+
"model = keras.Sequential([\n",
156+
" keras.layers.InputLayer(input_shape=(28, 28)),\n",
157+
" keras.layers.Reshape(target_shape=(28, 28, 1)),\n",
158+
" keras.layers.Conv2D(filters=12, kernel_size=(3, 3),\n",
158159
" activation=tf.nn.relu),\n",
159-
" tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
160-
" tf.keras.layers.Flatten(),\n",
161-
" tf.keras.layers.Dense(10)\n",
160+
" keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
161+
" keras.layers.Flatten(),\n",
162+
" keras.layers.Dense(10)\n",
162163
"])\n",
163164
"\n",
164165
"# Train the digit classification model\n",
165166
"model.compile(optimizer='adam',\n",
166-
" loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
167+
" loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
167168
" metrics=['accuracy'])\n",
168169
"\n",
169170
"model.fit(\n",
@@ -198,7 +199,7 @@
198199
"\n",
199200
"_, keras_file = tempfile.mkstemp('.h5')\n",
200201
"print('Saving model to: ', keras_file)\n",
201-
"tf.keras.models.save_model(model, keras_file, include_optimizer=False)"
202+
"keras.models.save_model(model, keras_file, include_optimizer=False)"
202203
]
203204
},
204205
{
@@ -259,10 +260,10 @@
259260
"clustered_model = cluster_weights(model, **clustering_params)\n",
260261
"\n",
261262
"# Use smaller learning rate for fine-tuning\n",
262-
"opt = tf.keras.optimizers.Adam(learning_rate=1e-5)\n",
263+
"opt = keras.optimizers.Adam(learning_rate=1e-5)\n",
263264
"\n",
264265
"clustered_model.compile(\n",
265-
" loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
266+
" loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
266267
" optimizer=opt,\n",
267268
" metrics=['accuracy'])\n",
268269
"\n",
@@ -323,7 +324,7 @@
323324
"def print_model_weight_clusters(model):\n",
324325
"\n",
325326
" for layer in model.layers:\n",
326-
" if isinstance(layer, tf.keras.layers.Wrapper):\n",
327+
" if isinstance(layer, keras.layers.Wrapper):\n",
327328
" weights = layer.trainable_weights\n",
328329
" else:\n",
329330
" weights = layer.weights\n",
@@ -414,7 +415,7 @@
414415
"qat_model = tfmot.quantization.keras.quantize_model(stripped_clustered_model)\n",
415416
"\n",
416417
"qat_model.compile(optimizer='adam',\n",
417-
" loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
418+
" loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
418419
" metrics=['accuracy'])\n",
419420
"print('Train qat model:')\n",
420421
"qat_model.fit(train_images, train_labels, batch_size=128, epochs=1, validation_split=0.1)\n",
@@ -427,7 +428,7 @@
427428
" tfmot.experimental.combine.Default8BitClusterPreserveQuantizeScheme())\n",
428429
"\n",
429430
"cqat_model.compile(optimizer='adam',\n",
430-
" loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
431+
" loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
431432
" metrics=['accuracy'])\n",
432433
"print('Train cqat model:')\n",
433434
"cqat_model.fit(train_images, train_labels, batch_size=128, epochs=1, validation_split=0.1)"

0 commit comments

Comments
 (0)