pqat_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 o pruning (PQAT)

Visão geral

Este é um exemplo completo que mostra o uso da API de treinamento consciente de quantização que preserva o pruning (PQAT), 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, usando a API de esparsidade, e verá a exatidão.

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

4. Aplicará o PQAT e observará se a esparsidade aplicada antes foi preservada.

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

6. Comparará a exatidão de modelo alcançada pelo PQAT 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 pruning

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

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

Aplique a API prune_low_magnitude() para podar 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 a exatidão. 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 pruning.

Defina o modelo e aplique a API de esparsidade

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

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

pruned_model.summary()

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 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:
            # ignore auxiliary quantization weights
            if "quantize_layer" 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})",
            )

Confira se o modelo foi podado corretamente. Precisamos retirar o wrapper de pruning primeiro.

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

print_model_weights_sparsity(stripped_pruned_model)

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

_, pruned_model_accuracy = pruned_model.evaluate(
  test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)
print('Pruned test accuracy:', pruned_model_accuracy)

Aplique o QAT e o PQAT e confira o efeito na esparsidade do modelo em ambos os casos

Em seguida, vamos aplicar o QAT e o QAT que preserva o pruning (PQAT) no modelo podado e observar se o PQAT preserva a esparsidade no modelo podado. Perceba que retiramos os wrappers de pruning do modelo com tfmot.sparsity.keras.strip_pruning antes de aplicar a API de PQAT.

# QAT
qat_model = tfmot.quantization.keras.quantize_model(stripped_pruned_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)

# PQAT
quant_aware_annotate_model = tfmot.quantization.keras.quantize_annotate_model(
              stripped_pruned_model)
pqat_model = tfmot.quantization.keras.quantize_apply(
              quant_aware_annotate_model,
              tfmot.experimental.combine.Default8BitPrunePreserveQuantizeScheme())

pqat_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
print('Train pqat model:')
pqat_model.fit(train_images, train_labels, batch_size=128, epochs=1, validation_split=0.1)

print("QAT Model sparsity:")
print_model_weights_sparsity(qat_model)
print("PQAT Model sparsity:")
print_model_weights_sparsity(pqat_model)

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

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

Como esse modelo é pequeno, a diferença entre os dois modelos não é muito perceptível. A aplicação do pruning e do PQAT 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)
    
# PQAT model
converter = tf.lite.TFLiteConverter.from_keras_model(pqat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
pqat_tflite_model = converter.convert()
pqat_model_file = 'pqat_model.tflite'
# Save the model.
with open(pqat_model_file, 'wb') as f:
    f.write(pqat_tflite_model)
    
print("QAT model size: ", get_gzipped_model_size(qat_model_file), ' KB')
print("PQAT model size: ", get_gzipped_model_size(pqat_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 e a quantização e veja se a exatidão do TensorFlow persiste no back-end do TFLite.

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

pqat_test_accuracy = eval_model(interpreter)

print('Pruned and quantized TFLite test_accuracy:', pqat_test_accuracy)
print('Pruned TF test accuracy:', pruned_model_accuracy)

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

Em seguida, vamos usar a quantização pós-treinamento (sem ajustes) no modelo podado e verificar a exatidão em relação ao modelo de PQAT. Isso demonstra por que você precisa usar o PQAT para melhorar a exatidão do modelo quantizado.

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 PQAT obtido anteriormente. Observe que o modelo quantizado com ajustes alcança maior exatidão.

converter = tf.lite.TFLiteConverter.from_keras_model(stripped_pruned_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('PQAT TFLite test_accuracy:', pqat_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, podá-lo usando a API de esparsidade e aplicar o treinamento consciente de quantização que preserva a esparsidade (PQAT) para preservar a esparsidade e usar o QAT. O modelo do PQAT final foi comparado ao do QAT para mostrar que a esparsidade é preservada no primeiro e perdida no último. Em seguida, os modelos foram convertidos para o TFLite, mostrando os benefícios da compressão ao usar as técnicas de pruning e PQAT 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 PQAT foi comparado a um modelo podado e quantizado com a API de quantização pós-treinamento, demonstrando a vantagem do PQAT em recuperar a perda de exatidão da quantização normal.