pqat_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 (PQAT)

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 poda (PQAT), 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, usando la API de dispersión, y comprobar la precisión.

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

4. Aplicar el PQAT y observar que se ha conservado la dispersión que se aplicó anteriormente.

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

6. Comparar la precisión que obtiene el modelo de PQAT 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 podar

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

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

Aplique la API prune_low_magnitude() para podar 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 poda.

Definir el modelo y aplicar la API de dispersión

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

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

pruned_model.summary()

Ajustar el modelo 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 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:
            # ignore auxiliary quantization weights
            if "quantize_layer" 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})",
            )

Verifique que se haya podado el modelo correctamente. Primero debemos quitar el contenedor de la poda.

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

print_model_weights_sparsity(stripped_pruned_model)

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

_, pruned_model_accuracy = pruned_model.evaluate(
  test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)
print('Pruned test accuracy:', pruned_model_accuracy)

Aplicar el QAT y el PQAT y verificar el efecto sobre la dispersión del modelo en ambos casos

A continuación, aplicamos el QAT y el QAT que preserva la poda (PQAT) en el modelo podado y observamos que el PQAT preserva la dispersión en su modelo podado. Tenga en cuenta que eliminamos los contenedores de poda de su modelo podado con tfmot.sparsity.keras.strip_pruning antes de aplicar la API de PQAT.

# QAT
qat_model = tfmot.quantization.keras.quantize_model(stripped_pruned_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)

# PQAT
quant_aware_annotate_model = tfmot.quantization.keras.quantize_annotate_model(
              stripped_pruned_model)
pqat_model = tfmot.quantization.keras.quantize_apply(
              quant_aware_annotate_model,
              tfmot.experimental.combine.Default8BitPrunePreserveQuantizeScheme())

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

print("QAT Model sparsity:")
print_model_weights_sparsity(qat_model)
print("PQAT Model sparsity:")
print_model_weights_sparsity(pqat_model)

Ver los beneficios de compresión del modelo de PQAT

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

Al tratarse de un modelo pequeño, la diferencia entre ambos modelos no se nota mucho. Aplicar la poda y el PQAT a un modelo de producción más grande producirí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)
    
# PQAT model
converter = tf.lite.TFLiteConverter.from_keras_model(pqat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
pqat_tflite_model = converter.convert()
pqat_model_file = 'pqat_model.tflite'
# Save the model.
with open(pqat_model_file, 'wb') as f:
    f.write(pqat_tflite_model)
    
print("QAT model size: ", get_gzipped_model_size(qat_model_file), ' KB')
print("PQAT model size: ", get_gzipped_model_size(pqat_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 y cuantizado, y verá que la precisión de TensorFlow persiste hasta el backend de TFLite.

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

pqat_test_accuracy = eval_model(interpreter)

print('Pruned and quantized TFLite test_accuracy:', pqat_test_accuracy)
print('Pruned TF test accuracy:', pruned_model_accuracy)

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

A continuación, usamos la cuantización normal posentrenamiento (sin ajuste) en el modelo podado y verificamos su precisión en comparación con el modelo de PQAT. Esto demuestra por qué debería usar el PQAT para mejorar la precisión del modelo cuantizado.

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 PQAT adquirido previamente. Tenga en cuenta que el modelo cuantizado con ajuste logra una mayor precisión.

converter = tf.lite.TFLiteConverter.from_keras_model(stripped_pruned_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('PQAT TFLite test_accuracy:', pqat_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, podarlo con la API de dispersión y aplicar el entrenamiento con reconocimiento de la cuantización que preserva la dispersión (PQAT) para preservar la dispersión al usar el QAT. El modelo de PQAT final se comparó con el de QAT para mostrar que la dispersión se conserva en el primero y se pierde en el segundo. Luego, los modelos se convirtieron a TFLite para mostrar los beneficios de compresión del encadenamiento de la poda y las técnicas de optimización del modelo de PQAT y se evaluó el modelo de TFLite para garantizar que la precisión persista en el backend de TFLite. Finalmente, el modelo de PQAT se comparó con un modelo podado cuantizado que se obtuvo con la API de cuantización posentrenamiento para demostrar la ventaja de PQAT en la recuperación de la pérdida de precisión de la cuantización normal.