<jupyter_start><jupyter_text>Run in Google Colab        View on GitHub    Weight Clustering Keras example  OverviewThis is an  example showing the usage of the **Weight Clustering** API, part of the TensorFlow Model Optimization Toolkit's collaborative optimization pipeline. We will use 4 and 8 clusters and see the affect on accuracy  SetupYou can run this Jupyter Notebook in local [virtualenv](https://www.tensorflow.org/install/pip?lang=python32.-create-a-virtual-environment-recommended) or [colab](https://colab.sandbox.google.com/).<jupyter_code>! pip install -q tensorflow-model-optimization<jupyter_output><empty_output><jupyter_text>Make the necessary import<jupyter_code>import tensorflow as tf

import numpy as np
import tempfile
import zipfile
import os<jupyter_output><empty_output><jupyter_text>Train a tf.keras model for MNIST without clustering1. load the dataset2. train and test images normalize3. Create Sequential model4. Compile the model with following parameters     * Use `adam` optimizer   * `SparseCategoricalCrossentropy`   * Optimize for `accuracy` metrics5. Run model.fit(..) with `train_images` and `train_labels` for 10 epochs and validation split of 0.1<jupyter_code># Load MNIST dataset
mnist = tf.keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# Normalize the input image so that each pixel value is between 0 to 1.
train_images = train_images / 255.0
test_images  = test_images / 255.0

model = tf.keras.Sequential([
  tf.keras.layers.InputLayer(input_shape=(28, 28)),
  tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
  tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3),
                         activation=tf.nn.relu),
  tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(10)
])

# Train the digit classification model
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.fit(
    train_images,
    train_labels,
    validation_split=0.1,
    epochs=1
)<jupyter_output>1688/1688 [==============================] - 19s 11ms/step - loss: 0.3035 - accuracy: 0.9139 - val_loss: 0.1229 - val_accuracy: 0.9677<jupyter_text>Evaluate the baseline model and save it for later usage<jupyter_code>_, baseline_model_accuracy = model.evaluate(
    test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)

_, keras_file = tempfile.mkstemp('.h5')
print('Saving model to: ', keras_file)
tf.keras.models.save_model(model, keras_file, include_optimizer=False)<jupyter_output>Baseline test accuracy: 0.958299994468689
Saving model to:  /tmp/tmp9gpn1bs4.h5<jupyter_text>Cluster and fine-tune the model with 4 and 8 clusters Apply the `cluster_weights()` API to cluster the whole pre-trained model to demonstrate and observe its effectiveness in reducing the model size when applying zip, while maintaining accuracy. For more details refer to the  [clustering comprehensive guide](https://www.tensorflow.org/model_optimization/guide/clustering/clustering_comprehensive_guide).  Define the model and apply the clustering API The model needs to be pre-trained before using the clustering API. This function wraps a keras model or layer with clustering functionality which clusters the layer's weights during training. For examples, using this with number_of_clusters equals 8 will ensure that each weight tensor has no more than 8 unique values.<jupyter_code>import tensorflow_model_optimization as tfmot

CentroidInitialization = tfmot.clustering.keras.CentroidInitialization
cluster_weights_4 = tfmot.clustering.keras.cluster_weights
model_4 = tf.keras.Sequential([
  tf.keras.layers.InputLayer(input_shape=(28, 28)),
  tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
  tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3),
                         activation=tf.nn.relu),
  tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(10)
])

# Train the digit classification model
model_4.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model_4.fit(
    train_images,
    train_labels,
    validation_split=0.1,
    epochs=1
)
cluster_weights_4_ = tfmot.clustering.keras.cluster_weights
clustering_params_4 = {
  'number_of_clusters': 4,
  'cluster_centroids_init': CentroidInitialization.KMEANS_PLUS_PLUS
}

clustered_model_4 = cluster_weights_4(model_4, **clustering_params_4)

# Use smaller learning rate for fine-tuning
opt = tf.keras.optimizers.Adam(learning_rate=1e-7)

clustered_model_4.compile(
  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
  optimizer=opt,
  metrics=['accuracy'])
cluster_weights_8 = tfmot.clustering.keras.cluster_weights
model_8 = tf.keras.Sequential([
  tf.keras.layers.InputLayer(input_shape=(28, 28)),
  tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
  tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3),
                         activation=tf.nn.relu),
  tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(10)
])

# Train the digit classification model
model_8.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model_8.fit(
    train_images,
    train_labels,
    validation_split=0.1,
    epochs=1
)
cluster_weights_8_ = tfmot.clustering.keras.cluster_weights
clustering_params_8 = {
  'number_of_clusters': 8,
  'cluster_centroids_init': CentroidInitialization.KMEANS_PLUS_PLUS
}

clustered_model_8 = cluster_weights_8(model_8, **clustering_params_4)

# Use smaller learning rate for fine-tuning
opt = tf.keras.optimizers.Adam(learning_rate=1e-7)

clustered_model_8.compile(
  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
  optimizer=opt,
  metrics=['accuracy'])<jupyter_output>1688/1688 [==============================] - 19s 11ms/step - loss: 0.3158 - accuracy: 0.9107 - val_loss: 0.1311 - val_accuracy: 0.9648<jupyter_text>Fine-tune the model and evaluate the accuracy against baseline Fine-tune the model with clustering for 3 epochs.<jupyter_code># Fine-tune model
clustered_model_4.fit(
  train_images,
  train_labels,
  epochs=3,
  validation_split=0.1)
# Fine-tune model
clustered_model_8.fit(
  train_images,
  train_labels,
  epochs=3,
  validation_split=0.1)<jupyter_output>Epoch 1/3
1688/1688 [==============================] - 24s 14ms/step - loss: 0.3037 - accuracy: 0.9223 - val_loss: 0.2528 - val_accuracy: 0.9383
Epoch 2/3
1688/1688 [==============================] - 22s 13ms/step - loss: 0.3012 - accuracy: 0.9231 - val_loss: 0.2504 - val_accuracy: 0.9387
Epoch 3/3
1688/1688 [==============================] - 23s 13ms/step - loss: 0.2978 - accuracy: 0.9238 - val_loss: 0.2471 - val_accuracy: 0.9388<jupyter_text>Define helper functions to calculate and print the number of clustering in each kernel of the model.<jupyter_code>def print_model_weight_clusters(model):

    for layer in model.layers:
        if isinstance(layer, tf.keras.layers.Wrapper):
            weights = layer.trainable_weights
        else:
            weights = layer.weights
        for weight in weights:
            # ignore auxiliary quantization weights
            if "quantize_layer" in weight.name:
                continue
            if "kernel" in weight.name:
                unique_count = len(np.unique(weight))
                print(
                    f"{layer.name}/{weight.name}: {unique_count} clusters "
                )<jupyter_output><empty_output><jupyter_text>Check that the model kernels were correctly clustered. We need to strip the clustering wrapper first.<jupyter_code>stripped_clustered_model_4 = tfmot.clustering.keras.strip_clustering(clustered_model_4)
print_model_weight_clusters(stripped_clustered_model_4)

stripped_clustered_model_8 = tfmot.clustering.keras.strip_clustering(clustered_model_8)
print_model_weight_clusters(stripped_clustered_model_8)<jupyter_output>conv2d_5/kernel:0: 4 clusters 
dense_5/kernel:0: 4 clusters 
conv2d_6/kernel:0: 4 clusters 
dense_6/kernel:0: 4 clusters<jupyter_text>For this example, there is minimal loss in test accuracy after clustering, compared to the baseline.<jupyter_code>_, clustered_model_accuracy_4 = clustered_model_4.evaluate(
  test_images, test_labels, verbose=0)
_, clustered_model_accuracy_8 = clustered_model_8.evaluate(
  test_images, test_labels, verbose=0)
print('Baseline test accuracy:', baseline_model_accuracy)
print('Clustered 4 clusters test accuracy:', clustered_model_accuracy_4)
print('Clustered 8 clusters test accuracy:', clustered_model_accuracy_8)<jupyter_output>Baseline test accuracy: 0.958299994468689
Clustered 4 clusters test accuracy: 0.9391999840736389
Clustered 8 clusters test accuracy: 0.9294999837875366
{'fraction_non_alphanumeric': 0.11168078466339723, 'fraction_numerical': 0.04458314757021846, 'mean_word_length': 4.860875244937949, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 31, 'https://': 3, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '422671', 'license': 'MIT', 'n_tokens_mistral': 3255, 'n_tokens_neox': 2943, 'n_words': 656, 'path': 'Chapter 09/weight_clustering_example_2_mnist.ipynb', 'repo_name': 'bpbpublications/Mastering-TensorFlow-2.x', 'size': '8972'}