pcqat_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 em TensorFlow.org

Executar no Google Colab

Ver no GitHub

Baixar notebook

Exemplo do Keras de treinamento consciente de quantização que preserva a esparsidade e os clusters (CQAT)

Visão geral

Este é um exemplo completo que mostra o uso da API de treinamento consciente de quantização que preserva a esparsidade e os clusters (PCQAT), parte do pipeline de otimização colaborativa do Kit de ferramentas para otimização de modelos do TensorFlow.

Outras páginas

Para uma introdução sobre o pipeline e outras técnicas disponíveis, confira a página de visão geral da otimização colaborativa.

Conteúdo

Neste tutorial, você:

1. Treinará um modelo tf.keras para o dataset MNIST do zero.

2. Ajustará o modelo com o pruning, verá a exatidão e observará se o pruning teve êxito.

3. Aplicará o clustering que preserva a esparsidade no modelo após o pruning e observará se a esparsidade aplicada antes foi preservada.

4. Aplicará o QAT e observará a perda de clusters.

5. Aplicará o CQAT e observará se a esparsidade e o clustering aplicados antes foram preservados.

6. Gerará um modelo do TFLite e observará os efeitos da aplicação do PCQAT nele.

7. Comparará os tamanhos de modelos diferentes para observar os benefícios da compressão ao aplicar a esparsidade, seguida por técnicas de otimização colaborativa: o clustering que preserva a esparsidade e o PCQAT.

8. Comparará a exatidão do modelo completamente otimizado com a exatidão do modelo de referência não otimizado.

Configuração

Você pode executar este Notebook do Jupyter no seu virtualenv local ou no colab. Para mais detalhes sobre como configurar as dependências, consulte o guia de instalação.

! pip install -q tensorflow-model-optimization

import tensorflow as tf

import numpy as np
import tempfile
import zipfile
import os

Treine um modelo tf.keras para o MNIST que passará pelo pruning e clustering

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

opt = tf.keras.optimizers.Adam(learning_rate=1e-3)

# Train the digit classification model
model.compile(optimizer=opt,
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.fit(
    train_images,
    train_labels,
    validation_split=0.1,
    epochs=10
)

Avalie o modelo de referência e salve-o para usar mais tarde

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

Faça o pruning e ajuste o modelo para 50% da esparsidade

Aplique a API prune_low_magnitude() para obter o modelo com pruning que será agrupado na próxima etapa. Consulte o guia completo de pruning para mais informações sobre a API de pruning.

Defina o modelo e aplique a API de esparsidade

Observe que é usado o modelo pré-treinado.

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

Ajuste o modelo e compare a exatidão dele com a referência

Ajuste o modelo com o pruning para 3 épocas.

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

Defina funções helper para calcular e imprimir a esparsidade e os clusters do 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:
            if "kernel" not in weight.name or "centroid" 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})",
            )

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

Vamos remover o wrapper de pruning primeiro e verificar se os kernels do modelo foram podados corretamente.

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

print_model_weights_sparsity(stripped_pruned_model)

Aplique o clustering que preserva a esparsidade e confira o efeito na esparsidade do modelo em ambos os casos

Em seguida, aplique o clustering que preserva a esparsidade no modelo após o pruning, observe o número de clusters e verifique se a esparsidade foi preservada.

import tensorflow_model_optimization as tfmot
from tensorflow_model_optimization.python.core.clustering.keras.experimental import (
    cluster,
)

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

cluster_weights = cluster.cluster_weights

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

sparsity_clustered_model = cluster_weights(stripped_pruned_model, **clustering_params)

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

print('Train sparsity preserving clustering model:')
sparsity_clustered_model.fit(train_images, train_labels,epochs=3, validation_split=0.1)

Remova o wrapper de clustering primeiro e verifique se o modelo foi podado e agrupado corretamente.

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

print("Model sparsity:\n")
print_model_weights_sparsity(stripped_clustered_model)

print("\nModel clusters:\n")
print_model_weight_clusters(stripped_clustered_model)

Aplique o QAT e o PCQAT e confira o efeito nos clusters e na esparsidade do modelo

Em seguida, aplicamos ambos o QAT e o PCQAT no modelo agrupado esparso e observamos se o PCQAT preserva a esparsidade e os clusters de peso no seu modelo. Observe que o modelo extraído é passado à API de QAT e PCQAT.

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

# PCQAT
quant_aware_annotate_model = tfmot.quantization.keras.quantize_annotate_model(
              stripped_clustered_model)
pcqat_model = tfmot.quantization.keras.quantize_apply(
              quant_aware_annotate_model,
              tfmot.experimental.combine.Default8BitClusterPreserveQuantizeScheme(preserve_sparsity=True))

pcqat_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
print('Train pcqat model:')
pcqat_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("\nQAT Model sparsity:")
print_model_weights_sparsity(qat_model)
print("\nPCQAT Model clusters:")
print_model_weight_clusters(pcqat_model)
print("\nPCQAT Model sparsity:")
print_model_weights_sparsity(pcqat_model)

Veja os benefícios da compressão do modelo de PCQAT

Defina a função helper para obter um arquivo de modelo compactado.

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

Observe se aplicar a esparsidade, o clustering e o PCQAT ao modelo gera benefícios de compressão significativos.

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

# PCQAT model
converter = tf.lite.TFLiteConverter.from_keras_model(pcqat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
pcqat_tflite_model = converter.convert()
pcqat_model_file = 'pcqat_model.tflite'
# Save the model.
with open(pcqat_model_file, 'wb') as f:
    f.write(pcqat_tflite_model)

print("QAT model size: ", get_gzipped_model_size(qat_model_file), ' KB')
print("PCQAT model size: ", get_gzipped_model_size(pcqat_model_file), ' KB')

Veja a persistência da exatidão do TF para o TFLite

Defina uma função helper para avaliar o modelo do TFLite com o dataset de teste.

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

Avalie o modelo após o pruning, o clustering e a quantização e veja se a exatidão do TensorFlow persiste no back-end do TFLite.

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

pcqat_test_accuracy = eval_model(interpreter)

print('Pruned, clustered and quantized TFLite test_accuracy:', pcqat_test_accuracy)
print('Baseline TF test accuracy:', baseline_model_accuracy)

Conclusão

Neste tutorial, você aprendeu a criar um modelo, fazer o pruning dele usando a API prune_low_magnitude() e aplicar o clustering que preserva a esparsidade usando a API cluster_weights() para preservar a esparsidade e agrupar os pesos.

Em seguida, o treinamento consciente de quantização que preserva a esparsidade e os clusters (PCQAT) foi aplicado para preservar a esparsidade e os clusters do modelo ao usar o QAT. O modelo PCQAT final foi comparado ao QAT para mostrar que a esparsidade e os clusters são preservados no primeiro e perdidos no último.

Em seguida, os modelos foram convertidos para o TFLite, mostrando os benefícios da compressão ao usar as técnicas de esparsidade, clustering e PCQAT em cadeia para a otimização do modelo, e o modelo do TFLite foi avaliado para garantir a persistência da eficácia no back-end do TFLite.

Por fim, a exatidão do modelo do TFLite de PCQAT foi comparada ao modelo de referência antes da otimização para mostrar que as técnicas de otimização colaborativa conseguiram obter benefícios de compressão e manter uma exatidão semelhante em comparação ao modelo original.