¡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

Guía completa de agrupación de pesos

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

Bienvenido a la guía completa para el agrupamiento de peso, que forma parte del conjunto de herramientas TensorFlow modelo de optimización.

Esta página documenta varios casos de uso y muestra cómo utilizar la API para cada uno. Una vez que sepas lo que las API que necesita, encontrar los parámetros y los detalles de bajo nivel en los documentos de la API :

En esta guía, se cubren los siguientes casos de uso:

  • Defina un modelo agrupado.
  • Controle y deserialice un modelo agrupado.
  • Mejore la precisión del modelo agrupado.
  • Solo para la implementación, debe seguir los pasos necesarios para ver los beneficios de la compresión.

Configuración

! pip install -q tensorflow-model-optimization

import tensorflow as tf
import numpy as np
import tempfile
import os
import tensorflow_model_optimization as tfmot

input_dim = 20
output_dim = 20
x_train = np.random.randn(1, input_dim).astype(np.float32)
y_train = tf.keras.utils.to_categorical(np.random.randn(1), num_classes=output_dim)

def setup_model():
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(input_dim, input_shape=[input_dim]),
      tf.keras.layers.Flatten()
  ])
  return model

def train_model(model):
  model.compile(
      loss=tf.keras.losses.categorical_crossentropy,
      optimizer='adam',
      metrics=['accuracy']
  )
  model.summary()
  model.fit(x_train, y_train)
  return model

def save_model_weights(model):
  _, pretrained_weights = tempfile.mkstemp('.h5')
  model.save_weights(pretrained_weights)
  return pretrained_weights

def setup_pretrained_weights():
  model= setup_model()
  model = train_model(model)
  pretrained_weights = save_model_weights(model)
  return pretrained_weights

def setup_pretrained_model():
  model = setup_model()
  pretrained_weights = setup_pretrained_weights()
  model.load_weights(pretrained_weights)
  return model

def save_model_file(model):
  _, keras_file = tempfile.mkstemp('.h5') 
  model.save(keras_file, include_optimizer=False)
  return keras_file

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

  keras_file = save_model_file(model)

  _, zipped_file = tempfile.mkstemp('.zip')
  with zipfile.ZipFile(zipped_file, 'w', compression=zipfile.ZIP_DEFLATED) as f:
    f.write(keras_file)
  return os.path.getsize(zipped_file)

setup_model()
pretrained_weights = setup_pretrained_weights()

Definir un modelo agrupado

Agrupar un modelo completo (secuencial y funcional)

Consejos para una mejor precisión del modelo:

  • Debe pasar un modelo previamente entrenado con una precisión aceptable a esta API. El entrenamiento de modelos desde cero con agrupamiento da como resultado una precisión insatisfactoria.
  • En algunos casos, agrupar determinadas capas tiene un efecto perjudicial sobre la precisión del modelo. Marque "Agrupar algunas capas" para ver cómo omitir la agrupación de las capas que más afectan la precisión.

Para agrupar todas las capas, aplicar tfmot.clustering.keras.cluster_weights al modelo.

import tensorflow_model_optimization as tfmot

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

clustering_params = {
  'number_of_clusters': 3,
  'cluster_centroids_init': CentroidInitialization.KMEANS_PLUS_PLUS
}

model = setup_model()
model.load_weights(pretrained_weights)

clustered_model = cluster_weights(model, **clustering_params)

clustered_model.summary()
Model: "sequential_2"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
cluster_dense_2 (ClusterWeig (None, 20)                823       
_________________________________________________________________
cluster_flatten_2 (ClusterWe (None, 20)                0         
=================================================================
Total params: 823
Trainable params: 423
Non-trainable params: 400
_________________________________________________________________

Agrupar algunas capas (modelos secuenciales y funcionales)

Consejos para una mejor precisión del modelo:

  • Debe pasar un modelo previamente entrenado con una precisión aceptable a esta API. El entrenamiento de modelos desde cero con agrupamiento da como resultado una precisión insatisfactoria.
  • Cluster capas posteriores con los parámetros más redundantes (por ejemplo tf.keras.layers.Dense , tf.keras.layers.Conv2D ), en oposición a las primeras capas.
  • Congele las primeras capas antes de las capas agrupadas durante el ajuste fino. Trate el número de capas congeladas como un hiperparámetro. Empíricamente, congelar la mayoría de las primeras capas es ideal para la API de agrupación en clúster actual.
  • Evite agrupar capas críticas (por ejemplo, mecanismo de atención).

Más: los tfmot.clustering.keras.cluster_weights docs API proporcionan detalles sobre cómo variar la configuración de la agrupación por capa.

# Create a base model
base_model = setup_model()
base_model.load_weights(pretrained_weights)

# Helper function uses `cluster_weights` to make only 
# the Dense layers train with clustering
def apply_clustering_to_dense(layer):
  if isinstance(layer, tf.keras.layers.Dense):
    return cluster_weights(layer, **clustering_params)
  return layer

# Use `tf.keras.models.clone_model` to apply `apply_clustering_to_dense` 
# to the layers of the model.
clustered_model = tf.keras.models.clone_model(
    base_model,
    clone_function=apply_clustering_to_dense,
)

clustered_model.summary()
Model: "sequential_3"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
cluster_dense_3 (ClusterWeig (None, 20)                823       
_________________________________________________________________
flatten_3 (Flatten)          (None, 20)                0         
=================================================================
Total params: 823
Trainable params: 423
Non-trainable params: 400
_________________________________________________________________

Agrupar la capa personalizada de Keras o especificar qué pesos de capa agrupar

tfmot.clustering.keras.ClusterableLayer sirve para dos casos de uso:

  1. Agrupe cualquier capa que no sea compatible de forma nativa, incluida una capa personalizada de Keras.
  2. Especifique qué pesos de una capa admitida se agruparán.

Para un ejemplo, los valores predeterminados de API sólo para agrupar el núcleo de la Dense capa. El siguiente ejemplo muestra cómo modificarlo para agrupar también el sesgo. Tenga en cuenta que cuando se deriva de la capa Keras, es necesario reemplazar la función get_clusterable_weights , donde se especifica el nombre de la variable entrenable a agruparse y de la propia variable de entrenable. Por ejemplo, si devuelve una lista vacía [], no se podrán agrupar pesos.

Error común: agrupar el sesgo por lo general perjudica la exactitud del modelo demasiado.

class MyDenseLayer(tf.keras.layers.Dense, tfmot.clustering.keras.ClusterableLayer):

  def get_clusterable_weights(self):
   # Cluster kernel and bias. This is just an example, clustering
   # bias usually hurts model accuracy.
   return [('kernel', self.kernel), ('bias', self.bias)]

# Use `cluster_weights` to make the `MyDenseLayer` layer train with clustering as usual.
model_for_clustering = tf.keras.Sequential([
  tfmot.clustering.keras.cluster_weights(MyDenseLayer(20, input_shape=[input_dim]), **clustering_params),
  tf.keras.layers.Flatten()
])

model_for_clustering.summary()
Model: "sequential_4"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
cluster_my_dense_layer (Clus (None, 20)                846       
_________________________________________________________________
flatten_4 (Flatten)          (None, 20)                0         
=================================================================
Total params: 846
Trainable params: 426
Non-trainable params: 420
_________________________________________________________________

También puede utilizar tfmot.clustering.keras.ClusterableLayer a agruparse una capa personalizada Keras. Para ello, se amplía tf.keras.Layer como de costumbre y aplicar el __init__ , call , y build funciones, pero también es necesario para extender la clusterable_layer.ClusterableLayer clase y poner en práctica:

  1. get_clusterable_weights , donde se especifica el kernel de peso a agruparse, como se muestra arriba.
  2. get_clusterable_algorithm , donde se especifica el algoritmo de agrupamiento para el tensor de peso. Esto se debe a que debe especificar cómo se forman los pesos de capa personalizados para la agrupación. La clase de algoritmo de agrupamiento devuelto debe ser derivada de la clustering_algorithm.ClusteringAlgorithm clase y la función get_pulling_indices debe ser sobrescrito. Un ejemplo de esta función, que soporta pesos de filas 1D, 2D, 3D y, se puede encontrar aquí .

Un ejemplo de este caso de uso se puede encontrar aquí .

Comprobar y deserializar un modelo agrupado

Su caso de uso: este código sólo es necesaria para el formato modelo HDF5 (no HDF5 pesos u otros formatos).

# Define the model.
base_model = setup_model()
base_model.load_weights(pretrained_weights)
clustered_model = cluster_weights(base_model, **clustering_params)

# Save or checkpoint the model.
_, keras_model_file = tempfile.mkstemp('.h5')
clustered_model.save(keras_model_file, include_optimizer=True)

# `cluster_scope` is needed for deserializing HDF5 models.
with tfmot.clustering.keras.cluster_scope():
  loaded_model = tf.keras.models.load_model(keras_model_file)

loaded_model.summary()
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.
Model: "sequential_5"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
cluster_dense_4 (ClusterWeig (None, 20)                823       
_________________________________________________________________
cluster_flatten_5 (ClusterWe (None, 20)                0         
=================================================================
Total params: 823
Trainable params: 423
Non-trainable params: 400
_________________________________________________________________

Mejorar la precisión del modelo agrupado

Para su caso de uso específico, hay algunos consejos que puede considerar:

  • La inicialización del centroide juega un papel clave en la precisión del modelo optimizado final. En general, la inicialización de kmeans ++ supera a la inicialización lineal, de densidad y aleatoria. Cuando no se usa kmeans ++, la inicialización lineal tiende a superar la densidad y la inicialización aleatoria, ya que no tiende a perder grandes pesos. Sin embargo, se ha observado que la inicialización de densidad proporciona una mayor precisión en el caso de utilizar muy pocos grupos en pesos con distribuciones bimodales.

  • Establezca una tasa de aprendizaje inferior a la utilizada en el entrenamiento al ajustar el modelo agrupado.

  • Para obtener ideas generales para mejorar la precisión del modelo, busque sugerencias para sus casos de uso en "Definir un modelo agrupado".

Despliegue

Exportar modelo con compresión de tamaño

Error común: ambos strip_clustering y la aplicación de un algoritmo de compresión estándar (por ejemplo, a través de gzip) son necesarias para ver las ventajas de la compresión de la agrupación.

model = setup_model()
clustered_model = cluster_weights(model, **clustering_params)

clustered_model.compile(
    loss=tf.keras.losses.categorical_crossentropy,
    optimizer='adam',
    metrics=['accuracy']
)

clustered_model.fit(
    x_train,
    y_train
)

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

print("final model")
final_model.summary()

print("\n")
print("Size of gzipped clustered model without stripping: %.2f bytes" 
      % (get_gzipped_model_size(clustered_model)))
print("Size of gzipped clustered model with stripping: %.2f bytes" 
      % (get_gzipped_model_size(final_model)))
1/1 [==============================] - 0s 343ms/step - loss: 16.1181 - accuracy: 0.0000e+00
final model
Model: "sequential_6"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense_5 (Dense)              (None, 20)                420       
_________________________________________________________________
flatten_6 (Flatten)          (None, 20)                0         
=================================================================
Total params: 420
Trainable params: 420
Non-trainable params: 0
_________________________________________________________________


Size of gzipped clustered model without stripping: 3447.00 bytes
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
Size of gzipped clustered model with stripping: 1433.00 bytes