cqat_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 os clusters (CQAT)

Visão geral

Este é um exemplo completo que mostra o uso da API de treinamento consciente de quantização que preserva os clusters (CQAT), 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 ao aplicar o clustering e verá a exatidão.

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

4. Aplicará o CQAT e observará se o clustering aplicado antes foi preservado.

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

6. Comparará a exatidão de modelo alcançada pelo CQAT com um modelo quantizado com a quantização pós-treinamento.

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 sem o 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)
])

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

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)

Agrupe e ajuste o modelo com 8 clusters

Aplique a API cluster_weights() para agrupar todo o modelo pré-treinado a fim de demonstrar e observar a eficácia em reduzir o tamanho do modelo após a aplicação do zip, mantendo uma exatidão adequada. Para saber como melhor usar a API para alcançar a melhor proporção de compressão e ainda manter a exatidão alvo, consulte o guia completo de clustering.

Defina o modelo e aplique a API de clustering

O modelo precisa ser pré-treinado antes de usar a API de clustering.

import tensorflow_model_optimization as tfmot

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

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

clustered_model = cluster_weights(model, **clustering_params)

# Use smaller learning rate for fine-tuning
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()

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

Ajuste o modelo com o clustering para 3 épocas.

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

Defina funções helper para calcular e imprimir o número de clustering em cada kernel do modelo.

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

Confira se os kernels do modelo foram agrupados corretamente. Precisamos retirar o wrapper de clustering primeiro.

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

print_model_weight_clusters(stripped_clustered_model)

Nesse exemplo, há uma perda mínima na exatidão do teste após o clustering, em comparação com a referência.

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

Aplique o QAT e o CQAT e confira o efeito nos clusters do modelo em ambos os casos

Em seguida, vamos aplicar o QAT e o QAT que preserva os clusters (CQAT) no modelo agrupado e observar se o CQAT preserva os clusters de peso no modelo agrupado. Perceba que retiramos os wrappers de clustering do modelo com tfmot.clustering.keras.strip_clustering antes de aplicar a API de CQAT.

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

# CQAT
quant_aware_annotate_model = tfmot.quantization.keras.quantize_annotate_model(
              stripped_clustered_model)
cqat_model = tfmot.quantization.keras.quantize_apply(
              quant_aware_annotate_model,
              tfmot.experimental.combine.Default8BitClusterPreserveQuantizeScheme())

cqat_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
print('Train cqat model:')
cqat_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("CQAT Model clusters:")
print_model_weight_clusters(cqat_model)

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

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

Esse é um modelo pequeno. A aplicação do clustering e do CQAT a um modelo de produção maior geraria uma compressão mais 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)
    
# CQAT model
converter = tf.lite.TFLiteConverter.from_keras_model(cqat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
cqat_tflite_model = converter.convert()
cqat_model_file = 'cqat_model.tflite'
# Save the model.
with open(cqat_model_file, 'wb') as f:
    f.write(cqat_tflite_model)
    
print("QAT model size: ", get_gzipped_model_size(qat_model_file), ' KB')
print("CQAT model size: ", get_gzipped_model_size(cqat_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 clustering e a quantização e veja se a exatidão do TensorFlow persiste no back-end do TFLite.

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

cqat_test_accuracy = eval_model(interpreter)

print('Clustered and quantized TFLite test_accuracy:', cqat_test_accuracy)
print('Clustered TF test accuracy:', clustered_model_accuracy)

Aplique a quantização pós-treinamento e compare com o modelo de CQAT

Em seguida, vamos usar a quantização pós-treinamento (sem ajustes) no modelo agrupado e verificar a exatidão em relação ao modelo de CQAT. Isso demonstra por que você precisa usar o CQAT para melhorar a exatidão do modelo quantizado. A diferença pode não ser muito visível, porque o modelo de MNIST é bastante pequeno e sobreparametrizado.

Primeiro, defina um gerador para o dataset de calibração a partir das primeiras 1000 imagens de treinamento.

def mnist_representative_data_gen():
  for image in train_images[:1000]:  
    image = np.expand_dims(image, axis=0).astype(np.float32)
    yield [image]

Faça a quantização do modelo e compare a exatidão com o modelo de CQAT obtido anteriormente. Observe que o modelo quantizado com ajustes alcança maior exatidão.

converter = tf.lite.TFLiteConverter.from_keras_model(stripped_clustered_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('CQAT TFLite test_accuracy:', cqat_test_accuracy)
print('Post-training (no fine-tuning) TF test accuracy:', post_training_test_accuracy)

Conclusão

Neste tutorial, você aprendeu a criar um modelo, agrupá-lo usando a API cluster_weights() e aplicar o treinamento consciente de quantização que preserva os clusters (CQAT) para preservar os clusters e usar o QAT. O modelo do CQAT final foi comparado ao do QAT para mostrar que 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 clustering e CQAT em cadeia para a otimização do modelo. O modelo do TFLite foi avaliado para garantir a persistência da exatidão no back-end do TFLite. Por fim, o modelo de CQAT foi comparado a um modelo agrupado e quantizado com a API de quantização pós-treinamento, demonstrando demonstrar a vantagem do CQAT em recuperar a perda de exatidão da quantização normal.