cqat_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 para el entrenamiento con reconocimiento de la cuantización que preserva los grupos (CQAT)

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 los grupos (CQAT), 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 agrupación y comprobar la precisión.

3. Aplicar el QAT y observar la pérdida de grupos.

4. Aplicar el CQAT y observar que se conserva la agrupación aplicada anteriormente.

5. Generar un modelo de TFLite y observar los efectos de la aplicación de CQAT.

6. Comparar la precisión que obtiene el modelo de CQAT con un modelo cuantizado mediante la cuantización posentrenamiento.

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 MNIST sin agrupación

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

Agrupar y ajustar el modelo con 8 grupos

Aplique la API cluster_weights() para agrupar todo el modelo preentrenado para demostrar y observar su efectividad en la reducción del tamaño del modelo al comprimirlo, mientras se mantiene la precisión. Para saber cuál es la mejor manera de usar la API para lograr la mejor tasa de compresión y al mismo tiempo mantener la precisión prevista, consulte la guía completa de agrupación en clústeres .

Definir el modelo y aplicar la API de agrupación

El modelo debe preentrenarse antes de usar la API de agrupación.

import tensorflow_model_optimization as tfmot

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

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

clustered_model = cluster_weights(model, **clustering_params)

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

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

clustered_model.summary()

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

Ajuste el modelo con agrupación durante 3 épocas.

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

Defina funciones ayudantes para calcular e imprimir el número de agrupaciones en cada núcleo del modelo.

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

Verifique que los núcleos del modelo estén agrupados correctamente. Primero debemos eliminar el contenedor de agrupación.

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

print_model_weight_clusters(stripped_clustered_model)

Para este ejemplo, hay una pérdida mínima en la precisión de la prueba después de la agrupación, en comparación con la línea de base.

_, clustered_model_accuracy = clustered_model.evaluate(
  test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)
print('Clustered test accuracy:', clustered_model_accuracy)

Aplicar el QAT y el CQAT y verificar el efecto en los grupos de modelos en ambos casos

A continuación, aplicamos tanto el QAT como el QAT que preserva los grupos (CQAT) en el modelo agrupado y observamos que el CQAT preserva los grupos de peso en su modelo agrupado. Tenga en cuenta que eliminamos los contenedores de agrupación del modelo con tfmot.clustering.keras.strip_clustering antes de aplicar la API de CQAT.

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

# CQAT
quant_aware_annotate_model = tfmot.quantization.keras.quantize_annotate_model(
              stripped_clustered_model)
cqat_model = tfmot.quantization.keras.quantize_apply(
              quant_aware_annotate_model,
              tfmot.experimental.combine.Default8BitClusterPreserveQuantizeScheme())

cqat_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
print('Train cqat model:')
cqat_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("CQAT Model clusters:")
print_model_weight_clusters(cqat_model)

Vea los beneficios de compresión del modelo CQAT

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

Tenga en cuenta que este es un modelo pequeño. La aplicación de agrupación y el CQAT a un modelo de producción más grande tendría una compresión más significativa.

# 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)
    
# CQAT model
converter = tf.lite.TFLiteConverter.from_keras_model(cqat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
cqat_tflite_model = converter.convert()
cqat_model_file = 'cqat_model.tflite'
# Save the model.
with open(cqat_model_file, 'wb') as f:
    f.write(cqat_tflite_model)
    
print("QAT model size: ", get_gzipped_model_size(qat_model_file), ' KB')
print("CQAT model size: ", get_gzipped_model_size(cqat_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 agrupado y cuantizado, y verá que la precisión de TensorFlow persiste hasta el backend de TFLite.

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

cqat_test_accuracy = eval_model(interpreter)

print('Clustered and quantized TFLite test_accuracy:', cqat_test_accuracy)
print('Clustered TF test accuracy:', clustered_model_accuracy)

Aplicar la cuantización posentrenamiento y comparar con el modelo de CQAT

A continuación, usamos la cuantización posentrenamiento (sin ajuste) en el modelo agrupado y verificamos su precisión en comparación con el modelo de CQAT. Esto demuestra por qué debería usar el CQAT para mejorar la precisión del modelo cuantizado. Es posible que la diferencia no sea muy visible, porque el modelo MNIST es bastante pequeño y está sobreparametrizado.

Primero, defina un generador para el conjunto de datos de calibración a partir de las primeras 1000 imágenes de entrenamiento.

def mnist_representative_data_gen():
  for image in train_images[:1000]:  
    image = np.expand_dims(image, axis=0).astype(np.float32)
    yield [image]

Cuantice el modelo y compare la precisión con el modelo de CQAT adquirido previamente. Tenga en cuenta que el modelo cuantizado con ajuste logra una mayor precisión.

converter = tf.lite.TFLiteConverter.from_keras_model(stripped_clustered_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = mnist_representative_data_gen
post_training_tflite_model = converter.convert()
post_training_model_file = 'post_training_model.tflite'
# Save the model.
with open(post_training_model_file, 'wb') as f:
    f.write(post_training_tflite_model)
    
# Compare accuracy
interpreter = tf.lite.Interpreter(post_training_model_file)
interpreter.allocate_tensors()

post_training_test_accuracy = eval_model(interpreter)

print('CQAT TFLite test_accuracy:', cqat_test_accuracy)
print('Post-training (no fine-tuning) TF test accuracy:', post_training_test_accuracy)

Conclusión

En este tutorial, aprendió cómo crear un modelo, agruparlo con la API cluster_weights() y aplicar el entrenamiento con reconocimiento de la cuantización que preserva los grupos (CQAT) para preservar los grupos al usar el QAT. El modelo de CQAT final se comparó con el de QAT para mostrar que los grupos se conservan en el primero y se pierden en el segundo. Luego, los modelos se convirtieron a TFLite para mostrar los beneficios de compresión del encadenamiento de agrupaciones y las técnicas de optimización del modelo de CQAT y se evaluó el modelo de TFLite para garantizar que la precisión persista en el backend de TFLite. Finalmente, el modelo de CQAT se comparó con un modelo agrupado cuantizado que se obtuvo con la API de cuantización posentrenamiento para demostrar la ventaja de CQAT en la recuperación de la pérdida de precisión de la cuantización normal.