¡El Día de la Comunidad de ML es el 9 de noviembre! Únase a nosotros para recibir actualizaciones de TensorFlow, JAX, y más Más información

Agrupación de peso en el ejemplo de Keras

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar cuaderno

Descripción general

Bienvenido al ejemplo de un extremo a otro para la agrupación en clústeres de ponderación , parte del kit de herramientas de optimización de modelos de TensorFlow.

Otras paginas

Para obtener una introducción a qué es la agrupación por peso y para determinar si debe usarla (incluido lo que es compatible), consulte la página de descripción general .

Para encontrar rápidamente las API que necesita para su caso de uso (más allá de agrupar por completo un modelo con 16 clústeres), consulte la guía completa .

Contenido

En el tutorial, podrá:

  1. Entrene un modelo tf.keras para el conjunto de datos MNIST desde cero.
  2. Ajuste el modelo aplicando la API de agrupación de ponderaciones y observe la precisión.
  3. Cree modelos TF y TFLite 6 veces más pequeños a partir de la agrupación.
  4. Cree un modelo TFLite 8 veces más pequeño combinando agrupación de peso y cuantificación posterior al entrenamiento.
  5. Vea la persistencia de la precisión de TF a TFLite.

Configuración

Puede ejecutar este Jupyter Notebook 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
from tensorflow import keras

import numpy as np
import tempfile
import zipfile
import os

Entrene un modelo tf.keras para MNIST sin clustering

# Load MNIST dataset
mnist = 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

# Define the model architecture.
model = keras.Sequential([
    keras.layers.InputLayer(input_shape=(28, 28)),
    keras.layers.Reshape(target_shape=(28, 28, 1)),
    keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation=tf.nn.relu),
    keras.layers.MaxPooling2D(pool_size=(2, 2)),
    keras.layers.Flatten(),
    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
)
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step
Epoch 1/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.3008 - accuracy: 0.9148 - val_loss: 0.1216 - val_accuracy: 0.9687
Epoch 2/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.1221 - accuracy: 0.9651 - val_loss: 0.0861 - val_accuracy: 0.9758
Epoch 3/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0897 - accuracy: 0.9741 - val_loss: 0.0710 - val_accuracy: 0.9802
Epoch 4/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0727 - accuracy: 0.9787 - val_loss: 0.0719 - val_accuracy: 0.9803
Epoch 5/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0631 - accuracy: 0.9808 - val_loss: 0.0657 - val_accuracy: 0.9822
Epoch 6/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0554 - accuracy: 0.9833 - val_loss: 0.0601 - val_accuracy: 0.9820
Epoch 7/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0489 - accuracy: 0.9855 - val_loss: 0.0647 - val_accuracy: 0.9805
Epoch 8/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0442 - accuracy: 0.9869 - val_loss: 0.0575 - val_accuracy: 0.9845
Epoch 9/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0403 - accuracy: 0.9875 - val_loss: 0.0596 - val_accuracy: 0.9820
Epoch 10/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0362 - accuracy: 0.9888 - val_loss: 0.0588 - val_accuracy: 0.9833
<tensorflow.python.keras.callbacks.History at 0x7f0e6f780a58>

Evalúe el modelo de línea de base y guárdelo para su uso posterior

_, 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)
Baseline test accuracy: 0.9785000085830688
Saving model to:  /tmp/tmpjo5b6jen.h5

Ajuste el modelo previamente entrenado con agrupación

Aplique la API cluster_weights() a un modelo completo previamente entrenado para demostrar su efectividad en la reducción del tamaño del modelo después de aplicar zip mientras mantiene una precisión decente. Para conocer la mejor manera de equilibrar la precisión y la tasa de compresión para su caso de uso, consulte el ejemplo por capa en la guía completa .

Definir el modelo y aplicar la API de agrupación

Antes de pasar el modelo a la API de agrupación en clústeres, asegúrese de que esté entrenado y muestre una precisión aceptable.

import tensorflow_model_optimization as tfmot

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

clustering_params = {
  'number_of_clusters': 16,
  'cluster_centroids_init': CentroidInitialization.LINEAR
}

# Cluster a whole model
clustered_model = cluster_weights(model, **clustering_params)

# Use smaller learning rate for fine-tuning clustered model
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()
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
cluster_reshape (ClusterWeig (None, 28, 28, 1)         0         
_________________________________________________________________
cluster_conv2d (ClusterWeigh (None, 26, 26, 12)        136       
_________________________________________________________________
cluster_max_pooling2d (Clust (None, 13, 13, 12)        0         
_________________________________________________________________
cluster_flatten (ClusterWeig (None, 2028)              0         
_________________________________________________________________
cluster_dense (ClusterWeight (None, 10)                20306     
=================================================================
Total params: 20,442
Trainable params: 54
Non-trainable params: 20,388
_________________________________________________________________

Ajuste el modelo y evalúe la precisión con respecto a la línea de base

Ajuste el modelo con agrupamiento durante 1 época.

# Fine-tune model
clustered_model.fit(
  train_images,
  train_labels,
  batch_size=500,
  epochs=1,
  validation_split=0.1)
108/108 [==============================] - 2s 16ms/step - loss: 0.0453 - accuracy: 0.9851 - val_loss: 0.0699 - val_accuracy: 0.9802
<tensorflow.python.keras.callbacks.History at 0x7f0e543ffeb8>

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)
Baseline test accuracy: 0.9785000085830688
Clustered test accuracy: 0.9746000170707703

Cree modelos 6 veces más pequeños a partir de la agrupación en clústeres

Tanto strip_clustering como aplicar un algoritmo de compresión estándar (por ejemplo, a través de gzip) son necesarios para ver los beneficios de la compresión de la agrupación en clústeres.

Primero, cree un modelo comprimible para TensorFlow. Aquí, strip_clustering elimina todas las variables (por ejemplo, tf.Variable para almacenar los centroides del clúster y los índices) que el clúster solo necesita durante el entrenamiento, lo que de otra manera se sumaría al tamaño del modelo durante la inferencia.

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

_, clustered_keras_file = tempfile.mkstemp('.h5')
print('Saving clustered model to: ', clustered_keras_file)
tf.keras.models.save_model(final_model, clustered_keras_file, 
                           include_optimizer=False)
Saving clustered model to:  /tmp/tmpo83fpb0m.h5

Luego, cree modelos comprimibles para TFLite. Puede convertir el modelo agrupado a un formato que se pueda ejecutar en su backend objetivo. TensorFlow Lite es un ejemplo que puede utilizar para implementar en dispositivos móviles.

clustered_tflite_file = '/tmp/clustered_mnist.tflite'
converter = tf.lite.TFLiteConverter.from_keras_model(final_model)
tflite_clustered_model = converter.convert()
with open(clustered_tflite_file, 'wb') as f:
  f.write(tflite_clustered_model)
print('Saved clustered TFLite model to:', clustered_tflite_file)
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/training/tracking/tracking.py:111: Model.state_updates (from tensorflow.python.keras.engine.training) is deprecated and will be removed in a future version.
Instructions for updating:
This property should not be used in TensorFlow 2.0, as updates are applied automatically.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/training/tracking/tracking.py:111: Layer.updates (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.
Instructions for updating:
This property should not be used in TensorFlow 2.0, as updates are applied automatically.
INFO:tensorflow:Assets written to: /tmp/tmp4gcxcvlh/assets
Saved clustered TFLite model to: /tmp/clustered_mnist.tflite

Defina una función auxiliar para comprimir realmente los modelos a través de gzip y medir el tamaño comprimido.

def get_gzipped_model_size(file):
  # It returns the size of the gzipped model in bytes.
  import os
  import zipfile

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

Compare y vea que los modelos son 6 veces más pequeños debido a la agrupación

print("Size of gzipped baseline Keras model: %.2f bytes" % (get_gzipped_model_size(keras_file)))
print("Size of gzipped clustered Keras model: %.2f bytes" % (get_gzipped_model_size(clustered_keras_file)))
print("Size of gzipped clustered TFlite model: %.2f bytes" % (get_gzipped_model_size(clustered_tflite_file)))
Size of gzipped baseline Keras model: 78047.00 bytes
Size of gzipped clustered Keras model: 12524.00 bytes
Size of gzipped clustered TFlite model: 12141.00 bytes

Crear un modelo TFLite 8x más pequeño de la combinación de la agrupación de peso y después de la formación de cuantización

Puede aplicar la cuantificación posterior al entrenamiento al modelo agrupado para obtener beneficios adicionales.

converter = tf.lite.TFLiteConverter.from_keras_model(final_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()

_, quantized_and_clustered_tflite_file = tempfile.mkstemp('.tflite')

with open(quantized_and_clustered_tflite_file, 'wb') as f:
  f.write(tflite_quant_model)

print('Saved quantized and clustered TFLite model to:', quantized_and_clustered_tflite_file)
print("Size of gzipped baseline Keras model: %.2f bytes" % (get_gzipped_model_size(keras_file)))
print("Size of gzipped clustered and quantized TFlite model: %.2f bytes" % (get_gzipped_model_size(quantized_and_clustered_tflite_file)))
INFO:tensorflow:Assets written to: /tmp/tmpt2flzp4s/assets
INFO:tensorflow:Assets written to: /tmp/tmpt2flzp4s/assets
Saved quantized and clustered TFLite model to: /tmp/tmpgu3loy72.tflite
Size of gzipped baseline Keras model: 78047.00 bytes
Size of gzipped clustered and quantized TFlite model: 9240.00 bytes

Vea la persistencia de la precisión de TF a TFLite

Defina una función auxiliar para evaluar el modelo 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('Evaluated on {n} results so far.'.format(n=i))
    # 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úa el modelo, que se ha agrupado y cuantificado, y luego ves que la precisión de TensorFlow persiste en el backend de TFLite.

interpreter = tf.lite.Interpreter(model_content=tflite_quant_model)
interpreter.allocate_tensors()

test_accuracy = eval_model(interpreter)

print('Clustered and quantized TFLite test_accuracy:', test_accuracy)
print('Clustered TF test accuracy:', clustered_model_accuracy)
Evaluated on 0 results so far.
Evaluated on 1000 results so far.
Evaluated on 2000 results so far.
Evaluated on 3000 results so far.
Evaluated on 4000 results so far.
Evaluated on 5000 results so far.
Evaluated on 6000 results so far.
Evaluated on 7000 results so far.
Evaluated on 8000 results so far.
Evaluated on 9000 results so far.


Clustered and quantized TFLite test_accuracy: 0.9746
Clustered TF test accuracy: 0.9746000170707703

Conclusión

En este instructivo, vio cómo crear modelos agrupados con la API del kit de herramientas de optimización de modelos de TensorFlow. Más específicamente, ha pasado por un ejemplo de extremo a extremo para crear un modelo 8 veces más pequeño para MNIST con una diferencia de precisión mínima. Le recomendamos que pruebe esta nueva capacidad, que puede ser particularmente importante para la implementación en entornos con recursos limitados.