comprehensive_guide.ipynb

Kernel: Python 3

Copyright 2020 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.

Guía completa de poda

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver código fuente en GitHub

Descargar el bloc de notas

Le damos la bienvenida a la guía completa para la poda de pesos de Keras.

En esta página, se documentan varios casos de uso y se muestra cómo usar la API para cada uno. Una vez que sepa qué API necesita, busque los parámetros y los detalles de bajo nivel en los documentos de API:

• Si quiere ver los beneficios de la poda y lo que se admite, consulte la descripción general.

• Para ver un ejemplo único de principio a fin, consulte el ejemplo de poda.

Se tratan los siguientes casos de uso:

• Definir y entrenar un modelo podado.

  • Secuencial y funcional.

  • model.fit y bucles de entrenamiento personalizados de Keras

• Establecer punto de verificación y deserializar un modelo podado.

• Implementar un modelo podado y ver los beneficios de la compresión.

Para configurar el algoritmo de poda, consulte los documentos de la API tfmot.sparsity.keras.prune_low_magnitude.

Preparación

Para encontrar las API que necesita y comprender los propósitos, puede ejecutar, pero omitir la lectura de esta sección.

! pip install -q tensorflow-model-optimization

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

%load_ext tensorboard

import tempfile

input_shape = [20]
x_train = np.random.randn(1, 20).astype(np.float32)
y_train = tf.keras.utils.to_categorical(np.random.randn(1), num_classes=20)

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

def setup_pretrained_weights():
  model = setup_model()

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

  model.fit(x_train, y_train)

  _, pretrained_weights = tempfile.mkstemp('.tf')

  model.save_weights(pretrained_weights)

  return pretrained_weights

def get_gzipped_model_size(model):
  # Returns size of gzipped model, in bytes.
  import os
  import zipfile

  _, keras_file = tempfile.mkstemp('.h5')
  model.save(keras_file, include_optimizer=False)

  _, 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 el modelo

Podar un modelo completo (secuencial y funcional)

Consejos para una mejor precisión del modelo:

• Pruebe "Prune some layers" (podar algunas capas) para omitir la poda de las capas que reducen más la precisión.

• Por lo general, es mejor realizar ajustes con la poda que con el entrenamiento desde cero.

Para hacer que todo el modelo se entrene con poda, aplique tfmot.sparsity.keras.prune_low_magnitude al modelo.

base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended.

model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model)

model_for_pruning.summary()

Podar algunas capas (secuencial y funcional)

Podar un modelo puede tener un efecto negativo en la precisión. Puede podar selectivamente capas de un modelo para explorar el equilibrio entre precisión, velocidad y tamaño del modelo.

Consejos para una mejor precisión del modelo:

• Por lo general, es mejor realizar ajustes con la poda que con el entrenamiento desde cero.

• Pruebe podar las capas posteriores en lugar de las primeras.

• Evite podar capas críticas (por ejemplo, mecanismo de atención).

Más:

• En los documentos de la API tfmot.sparsity.keras.prune_low_magnitude se pueden encontrar detalles sobre cómo variar la configuración de la poda por capas.

En el siguiente ejemplo, pode solo las capas Dense.

# Create a base model
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy

# Helper function uses `prune_low_magnitude` to make only the 
# Dense layers train with pruning.
def apply_pruning_to_dense(layer):
  if isinstance(layer, tf.keras.layers.Dense):
    return tfmot.sparsity.keras.prune_low_magnitude(layer)
  return layer

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

model_for_pruning.summary()

Si bien en este ejemplo se usó el tipo de capa para decidir qué podar, la forma más fácil de podar una capa en particular es establecer su propiedad name y buscar ese nombre en clone_function.

print(base_model.layers[0].name)

Precisión del modelo más legible pero potencialmente menor

Esto no es compatible con el ajuste con poda, por lo que puede ser menos preciso que los ejemplos anteriores que admiten el ajuste.

Si bien prune_low_magnitude se puede aplicar al definir el modelo inicial, cargar los pesos después no funciona en los ejemplos siguientes.

Ejemplo funcional

# Use `prune_low_magnitude` to make the `Dense` layer train with pruning.
i = tf.keras.Input(shape=(20,))
x = tfmot.sparsity.keras.prune_low_magnitude(tf.keras.layers.Dense(10))(i)
o = tf.keras.layers.Flatten()(x)
model_for_pruning = tf.keras.Model(inputs=i, outputs=o)

model_for_pruning.summary()

Ejemplo secuencial

# Use `prune_low_magnitude` to make the `Dense` layer train with pruning.
model_for_pruning = tf.keras.Sequential([
  tfmot.sparsity.keras.prune_low_magnitude(tf.keras.layers.Dense(20, input_shape=input_shape)),
  tf.keras.layers.Flatten()
])

model_for_pruning.summary()

Podar la capa de Keras personalizada o modificar qué partes de la capa podar

Error común: podar el sesgo suele perjudicar demasiado la precisión del modelo.

tfmot.sparsity.keras.PrunableLayer sirve para dos casos de uso:

1. Podar una capa de Keras personalizada

2. Modificar partes de una capa de Keras integrada para podarla.

Por ejemplo, la API por defecto solo poda el núcleo de la capa Dense. En el siguiente ejemplo también se elimina el sesgo.

class MyDenseLayer(tf.keras.layers.Dense, tfmot.sparsity.keras.PrunableLayer):

  def get_prunable_weights(self):
    # Prune bias also, though that usually harms model accuracy too much.
    return [self.kernel, self.bias]

# Use `prune_low_magnitude` to make the `MyDenseLayer` layer train with pruning.
model_for_pruning = tf.keras.Sequential([
  tfmot.sparsity.keras.prune_low_magnitude(MyDenseLayer(20, input_shape=input_shape)),
  tf.keras.layers.Flatten()
])

model_for_pruning.summary()

Entrenar el modelo

Model.fit

Llame a la retrollamada tfmot.sparsity.keras.UpdatePruningStep durante el entrenamiento.

Para ayudar a depurar el entrenamiento, use la retrollamada tfmot.sparsity.keras.PruningSummaries.

# Define the model.
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model)

log_dir = tempfile.mkdtemp()
callbacks = [
    tfmot.sparsity.keras.UpdatePruningStep(),
    # Log sparsity and other metrics in Tensorboard.
    tfmot.sparsity.keras.PruningSummaries(log_dir=log_dir)
]

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

model_for_pruning.fit(
    x_train,
    y_train,
    callbacks=callbacks,
    epochs=2,
)

#docs_infra: no_execute
%tensorboard --logdir={log_dir}

Los usuarios que no son de Colab pueden ver los resultados de una ejecución anterior de este bloque de código en TensorBoard.dev.

Bucle de entrenamiento personalizado

Llame a la retrollamada tfmot.sparsity.keras.UpdatePruningStep durante el entrenamiento.

Para ayudar a depurar el entrenamiento, use la retrollamada tfmot.sparsity.keras.PruningSummaries.

# Define the model.
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model)

# Boilerplate
loss = tf.keras.losses.categorical_crossentropy
optimizer = tf.keras.optimizers.Adam()
log_dir = tempfile.mkdtemp()
unused_arg = -1
epochs = 2
batches = 1 # example is hardcoded so that the number of batches cannot change.

# Non-boilerplate.
model_for_pruning.optimizer = optimizer
step_callback = tfmot.sparsity.keras.UpdatePruningStep()
step_callback.set_model(model_for_pruning)
log_callback = tfmot.sparsity.keras.PruningSummaries(log_dir=log_dir) # Log sparsity and other metrics in Tensorboard.
log_callback.set_model(model_for_pruning)

step_callback.on_train_begin() # run pruning callback
for _ in range(epochs):
  log_callback.on_epoch_begin(epoch=unused_arg) # run pruning callback
  for _ in range(batches):
    step_callback.on_train_batch_begin(batch=unused_arg) # run pruning callback

    with tf.GradientTape() as tape:
      logits = model_for_pruning(x_train, training=True)
      loss_value = loss(y_train, logits)
      grads = tape.gradient(loss_value, model_for_pruning.trainable_variables)
      optimizer.apply_gradients(zip(grads, model_for_pruning.trainable_variables))

  step_callback.on_epoch_end(batch=unused_arg) # run pruning callback

#docs_infra: no_execute
%tensorboard --logdir={log_dir}

Los usuarios que no son de Colab pueden ver los resultados de una ejecución anterior de este bloque de código en TensorBoard.dev.

Mejorar la precisión del modelo podado

Primero, consulte los documentos de la API tfmot.sparsity.keras.prune_low_magnitude para comprender qué es un programa de poda y las matemáticas de cada tipo de programación de poda.

Consejos:

• Tenga una tasa de aprendizaje que no sea demasiado alta ni demasiado baja cuando el modelo esté podando. Considere la programación de poda como un hiperparámetro.

• Para probarlo rápidamente, intente experimentar podando un modelo hasta la presición final al comienzo del entrenamiento configurando begin_step en 0 con una programación de tfmot.sparsity.keras.ConstantSparsity. Es posible que obtengas buenos resultados.

• No podes con mucha frecuencia para darle tiempo al modelo a recuperarse. La programación de poda proporciona una frecuencia predeterminada aceptable.

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

Establecer punto de verificación y deserializar

Debe conservar el paso del optimizador durante el establecimiento de puntos de verificación. Esto significa que, si bien puede usar los modelos de Keras HDF5 para puntos de verificación, no puede usar pesos de Keras HDF5.

# Define the model.
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model)

_, keras_model_file = tempfile.mkstemp('.h5')

# Checkpoint: saving the optimizer is necessary (include_optimizer=True is the default).
model_for_pruning.save(keras_model_file, include_optimizer=True)

Lo anterior se aplica en general. El siguiente código solo es necesario para el formato del modelo HDF5 (no para los pesos HDF5 ni otros formatos).

# Deserialize model.
with tfmot.sparsity.keras.prune_scope():
  loaded_model = tf.keras.models.load_model(keras_model_file)

loaded_model.summary()

Implementar modelo podado

Exportar modelo con compresión de tamaño

Error común: tanto strip_pruning como la aplicación de un algoritmo de compresión estándar (por ejemplo, mediante gzip) son necesarios para ver los beneficios de la compresión de la poda.

# Define the model.
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model)

# Typically you train the model here.

model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)

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

print("\n")
print("Size of gzipped pruned model without stripping: %.2f bytes" % (get_gzipped_model_size(model_for_pruning)))
print("Size of gzipped pruned model with stripping: %.2f bytes" % (get_gzipped_model_size(model_for_export)))

Optimizaciones específicas de hardware

Una vez que diferentes backends permiten la poda para mejorar la latencia, el uso de la dispersión de bloques puede mejorar la latencia para cierto hardware.

Aumentar el tamaño del bloque disminuirá la dispersión máxima que se puede lograr para la precisión del modelo previsto. A pesar de esto, la latencia aún puede mejorar.

Para obtener detalles sobre lo que se admite para la dispersión de bloques, consulte los documentos de la API tfmot.sparsity.keras.prune_low_magnitude.

base_model = setup_model()

# For using intrinsics on a CPU with 128-bit registers, together with 8-bit
# quantized weights, a 1x16 block size is nice because the block perfectly
# fits into the register.
pruning_params = {'block_size': [1, 16]}
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model, **pruning_params)

model_for_pruning.summary()