pcqat_example.ipynb

Kernel: Python 3

Copyright 2021 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver en GitHub

Descargar el bloc de notas

Ejemplo de Keras de entrenamiento con reconocimiento de la cuantización que preserva la dispersión y los grupos (PCQAT)

Descripción general

Este es un ejemplo de principio a fin que muestra el uso de la API de entrenamiento con reconocimiento de la cuantización que preserva la dispersión y los grupos (PCQAT), parte del proceso de optimización colaborativa del kit de herramientas de optimización de modelos de TensorFlow.

Otras paginas

Para ver una introducción a la canalización y otras técnicas disponibles, consulte la página de descripción general de optimización colaborativa.

Contenido

En este tutorial podrá:

1. Entrenar un modelo tf.keras para el conjunto de datos MNIST desde cero.

2. Ajustar el modelo con poda, ver la precisión y observar que el modelo se poda con éxito.

3. Aplicar la dispersión que preserva la agrupación en el modelo podado y observar que la dispersión aplicada anteriormente se ha conservado.

4. Aplicar el QAT y observar la pérdida de dispersión y grupos.

5. Aplicar el PCQAT y observar que tanto la dispersión como la agrupación aplicada anteriormente se han conservado.

6. Generar un modelo de TFLite y observar los efectos de la aplicación de PCQAT.

7. Comparar los tamaños de los diferentes modelos para observar los beneficios de la compresión al aplicar la dispersión seguida de las técnicas de optimización colaborativa de la agrupación que preserva la dispersión y el PCQAT.

8. Comparar la precisión del modelo totalmente optimizado con la precisión del modelo de referencia no optimizado.

Preparación

Puede ejecutar este cuaderno Jupyter en su virtualenv o colab local. Para obtener detalles sobre la configuración de dependencias, consulte la guía de instalación.

! pip install -q tensorflow-model-optimization

import tensorflow as tf

import numpy as np
import tempfile
import zipfile
import os

Entrenar un modelo tf.keras para que MNIST pueda podarse y agruparse

# 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)
])

opt = tf.keras.optimizers.Adam(learning_rate=1e-3)

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

model.fit(
    train_images,
    train_labels,
    validation_split=0.1,
    epochs=10
)

Evaluar el modelo previsto y guardarlo para usarlo después

_, 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)

Podar y ajustar el modelo hasta un 50 % de dispersión

Aplique la API prune_low_magnitude() para obtenere el modelo podado que se agrupará en el siguiente paso. Consulte la guía completa de poda para obtener más información sobre la API de poda.

Definir el modelo y aplicar la API de dispersión

Tenga en cuenta que se usa el modelo preentrenado.

import tensorflow_model_optimization as tfmot

prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

pruning_params = {
      'pruning_schedule': tfmot.sparsity.keras.ConstantSparsity(0.5, begin_step=0, frequency=100)
  }

callbacks = [
  tfmot.sparsity.keras.UpdatePruningStep()
]

pruned_model = prune_low_magnitude(model, **pruning_params)

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

pruned_model.compile(
  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
  optimizer=opt,
  metrics=['accuracy'])

Ajustar el modelo, verificar la dispersión y evaluar la precisión con respecto a la línea de base.

Ajuste el modelo con poda durante 3 épocas.

# Fine-tune model
pruned_model.fit(
  train_images,
  train_labels,
  epochs=3,
  validation_split=0.1,
  callbacks=callbacks)

Defina funciones ayudantes para calcular e imprimir la dispersión y los grupos del modelo.

def print_model_weights_sparsity(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:
            if "kernel" not in weight.name or "centroid" in weight.name:
                continue
            weight_size = weight.numpy().size
            zero_num = np.count_nonzero(weight == 0)
            print(
                f"{weight.name}: {zero_num/weight_size:.2%} sparsity ",
                f"({zero_num}/{weight_size})",
            )

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 "
                )

Primero quitemos el contenedor de poda y luego verifiquemos que los núcleos del modelo se podaron correctamente.

stripped_pruned_model = tfmot.sparsity.keras.strip_pruning(pruned_model)

print_model_weights_sparsity(stripped_pruned_model)

Aplicar la dispersión que preserva la agrupación y comprobar su efecto en la dispersión del modelo en ambos casos

A continuación, aplique la dispersión que preserva la agrupación en el modelo podado y observe la cantidad de grupos y que se preserve la dispersión.

import tensorflow_model_optimization as tfmot
from tensorflow_model_optimization.python.core.clustering.keras.experimental import (
    cluster,
)

cluster_weights = tfmot.clustering.keras.cluster_weights
CentroidInitialization = tfmot.clustering.keras.CentroidInitialization

cluster_weights = cluster.cluster_weights

clustering_params = {
  'number_of_clusters': 8,
  'cluster_centroids_init': CentroidInitialization.KMEANS_PLUS_PLUS,
  'preserve_sparsity': True
}

sparsity_clustered_model = cluster_weights(stripped_pruned_model, **clustering_params)

sparsity_clustered_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

print('Train sparsity preserving clustering model:')
sparsity_clustered_model.fit(train_images, train_labels,epochs=3, validation_split=0.1)

Primero, quite el contenedor de agrupación y luego verifique que se haya podado y agrupado el modelo correctamente.

stripped_clustered_model = tfmot.clustering.keras.strip_clustering(sparsity_clustered_model)

print("Model sparsity:\n")
print_model_weights_sparsity(stripped_clustered_model)

print("\nModel clusters:\n")
print_model_weight_clusters(stripped_clustered_model)

Aplicar el QAT y el PCQAT y verificar el efecto en los grupos y la dispersión del modelo

A continuación, aplique el QAT y el PCQAT en el modelo agrupado y disperso y observe que el PCQAT preserva la dispersión de pesos y los grupos en su modelo. Tenga en cuenta que el modelo eliminado se pasa a la API de QAT y PCQAT.

# QAT
qat_model = tfmot.quantization.keras.quantize_model(stripped_clustered_model)

qat_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
print('Train qat model:')
qat_model.fit(train_images, train_labels, batch_size=128, epochs=1, validation_split=0.1)

# PCQAT
quant_aware_annotate_model = tfmot.quantization.keras.quantize_annotate_model(
              stripped_clustered_model)
pcqat_model = tfmot.quantization.keras.quantize_apply(
              quant_aware_annotate_model,
              tfmot.experimental.combine.Default8BitClusterPreserveQuantizeScheme(preserve_sparsity=True))

pcqat_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
print('Train pcqat model:')
pcqat_model.fit(train_images, train_labels, batch_size=128, epochs=1, validation_split=0.1)

print("QAT Model clusters:")
print_model_weight_clusters(qat_model)
print("\nQAT Model sparsity:")
print_model_weights_sparsity(qat_model)
print("\nPCQAT Model clusters:")
print_model_weight_clusters(pcqat_model)
print("\nPCQAT Model sparsity:")
print_model_weights_sparsity(pcqat_model)

Ver los beneficios de compresión del modelo de PCQAT

Defina la función ayudante para obtener el archivo del modelo comprimido.

def get_gzipped_model_size(file):
  # It returns the size of the gzipped model in kilobytes.

  _, zipped_file = tempfile.mkstemp('.zip')
  with zipfile.ZipFile(zipped_file, 'w', compression=zipfile.ZIP_DEFLATED) as f:
    f.write(file)

  return os.path.getsize(zipped_file)/1000

Observe cómo la aplicación de dispersión, agrupación y PCQAT a un modelo produce importantes beneficios de compresión.

# QAT model
converter = tf.lite.TFLiteConverter.from_keras_model(qat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
qat_tflite_model = converter.convert()
qat_model_file = 'qat_model.tflite'
# Save the model.
with open(qat_model_file, 'wb') as f:
    f.write(qat_tflite_model)

# PCQAT model
converter = tf.lite.TFLiteConverter.from_keras_model(pcqat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
pcqat_tflite_model = converter.convert()
pcqat_model_file = 'pcqat_model.tflite'
# Save the model.
with open(pcqat_model_file, 'wb') as f:
    f.write(pcqat_tflite_model)

print("QAT model size: ", get_gzipped_model_size(qat_model_file), ' KB')
print("PCQAT model size: ", get_gzipped_model_size(pcqat_model_file), ' KB')

Ver la persistencia de la precisión desde TF a TFLite

Defina una función ayudante para evaluar el modelo de TFLite en el conjunto de datos de prueba.

def eval_model(interpreter):
  input_index = interpreter.get_input_details()[0]["index"]
  output_index = interpreter.get_output_details()[0]["index"]

  # Run predictions on every image in the "test" dataset.
  prediction_digits = []
  for i, test_image in enumerate(test_images):
    if i % 1000 == 0:
      print(f"Evaluated on {i} results so far.")
    # Pre-processing: add batch dimension and convert to float32 to match with
    # the model's input data format.
    test_image = np.expand_dims(test_image, axis=0).astype(np.float32)
    interpreter.set_tensor(input_index, test_image)

    # Run inference.
    interpreter.invoke()

    # Post-processing: remove batch dimension and find the digit with highest
    # probability.
    output = interpreter.tensor(output_index)
    digit = np.argmax(output()[0])
    prediction_digits.append(digit)

  print('\n')
  # Compare prediction results with ground truth labels to calculate accuracy.
  prediction_digits = np.array(prediction_digits)
  accuracy = (prediction_digits == test_labels).mean()
  return accuracy

Evalúe el modelo, que ha sido podado, agrupado y cuantizado, y verá que la precisión de TensorFlow persiste hasta el backend de TFLite.

interpreter = tf.lite.Interpreter(pcqat_model_file)
interpreter.allocate_tensors()

pcqat_test_accuracy = eval_model(interpreter)

print('Pruned, clustered and quantized TFLite test_accuracy:', pcqat_test_accuracy)
print('Baseline TF test accuracy:', baseline_model_accuracy)

Conclusión

En este tutorial, aprendió cómo crear un modelo, podarlo con la API prune_low_magnitude() y aplicar la agrupación que preserva la dispersión con la API cluster_weights() para preservar la dispersión mientras se agrupan los pesos.

Luego, se aplicó el entrenamiento con reconocimiento de la cuantización para preservar la dispersión y los grupos (PCQAT) para preservar la dispersión del modelo y los grupos al usar el QAT. El modelo de PCQAT final se comparó con el de QAT para mostrar que la dispersión y los grupos se conservan en el primero y se pierden en el segundo.

Después, los modelos se convirtieron a TFLite para mostrar los beneficios de compresión del encadenamiento de dispersión, agrupación y técnicas de optimización de modelos de PCQAT y se evaluó el modelo de TFLite para garantizar que la precisión persista en el backend de TFLite.

Por último, se comparó la precisión del modelo de PCQAT de TFLite con la precisión del modelo de referencia previo a la optimización para mostrar que las técnicas de optimización colaborativa pudieron lograr los beneficios de la compresión y mantener una precisión similar en comparación con el modelo original.