training_example.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.

Exemplo de treinamento consciente de quantização no Keras

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Visão geral

Bem-vindo ao exemplo completo de treinamento consciente de quantização.

Outras páginas

Para uma introdução sobre o que é o treinamento consciente de quantização e para determinar se você deve usá-lo (incluindo a compatibilidade dele), veja a página de visão geral.

Para encontrar rapidamente as APIs necessárias para seu caso de uso (além de quantizar um modelo inteiro de 8 bits), veja o guia completo.

Resumo

Neste tutorial, você:

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

2. Ajustará o modelo ao aplicar a API de treinamento consciente de quantização, verá a exatidão e exportará um modelo consciente de quantização.

3. Usará o modelo para criar um modelo realmente quantizado para o back-end do TFLite.

4. Verá a persistência da exatidão no TFLite e um modelo 4x menor. Para conferir os benefícios da latência em dispositivos móveis, teste os exemplos no repositório de apps do TFLite.

Configuração

! pip install -q tensorflow
! pip install -q tensorflow-model-optimization

import tempfile
import os

import tensorflow as tf

from tensorflow import keras

Treine um modelo para o MNIST sem o treinamento consciente de quantização

# Load MNIST dataset
mnist = 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

# Define the model architecture.
model = keras.Sequential([
  keras.layers.InputLayer(input_shape=(28, 28)),
  keras.layers.Reshape(target_shape=(28, 28, 1)),
  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),
  keras.layers.MaxPooling2D(pool_size=(2, 2)),
  keras.layers.Flatten(),
  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,
  epochs=1,
  validation_split=0.1,
)

Clone e ajuste o modelo pré-treinado com o treinamento consciente de quantização

Defina o modelo

Aplique o treinamento consciente de quantização ao modelo inteiro e veja isso no resumo do modelo. Todas as camadas agora têm o prefixo "quant".

Observe que o modelo resultante é consciente de quantização, mas não quantizado (por exemplo, os pesos são float32 em vez de int8). As seções seguintes mostram como criar um modelo quantizado a partir do consciente de quantização.

No guia completo, você pode ver como quantizar algumas camadas para melhorar a exatidão do modelo.

import tensorflow_model_optimization as tfmot

quantize_model = tfmot.quantization.keras.quantize_model

# q_aware stands for for quantization aware.
q_aware_model = quantize_model(model)

# `quantize_model` requires a recompile.
q_aware_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

q_aware_model.summary()

Treine e compare o modelo com a referência

Para mostrar os ajustes após o treinamento do modelo por apenas uma época, use o treinamento consciente de quantização em um subconjunto dos dados de treinamento.

train_images_subset = train_images[0:1000] # out of 60000
train_labels_subset = train_labels[0:1000]

q_aware_model.fit(train_images_subset, train_labels_subset,
                  batch_size=500, epochs=1, validation_split=0.1)

Nesse exemplo, há pouca ou nenhuma perda na exatidão do teste após o treinamento consciente de quantização, em comparação com a referência.

_, baseline_model_accuracy = model.evaluate(
    test_images, test_labels, verbose=0)

_, q_aware_model_accuracy = q_aware_model.evaluate(
   test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)
print('Quant test accuracy:', q_aware_model_accuracy)

Crie um modelo quantizado para o back-end do TFLite

Depois disso, você realmente tem um modelo quantizado com pesos int8 e ativações uint8.

converter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

quantized_tflite_model = converter.convert()

Verá 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.

import numpy as np

def evaluate_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('Evaluated on {n} results so far.'.format(n=i))
    # 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 quantizado e veja se a exatidão do TensorFlow persiste no back-end do TFLite.

interpreter = tf.lite.Interpreter(model_content=quantized_tflite_model)
interpreter.allocate_tensors()

test_accuracy = evaluate_model(interpreter)

print('Quant TFLite test_accuracy:', test_accuracy)
print('Quant TF test accuracy:', q_aware_model_accuracy)

Veja se o modelo está 4x menor com a quantização

Crie um modelo do TFLite float e veja se o modelo do TFLite quantizado está 4x menor.

# Create float TFLite model.
float_converter = tf.lite.TFLiteConverter.from_keras_model(model)
float_tflite_model = float_converter.convert()

# Measure sizes of models.
_, float_file = tempfile.mkstemp('.tflite')
_, quant_file = tempfile.mkstemp('.tflite')

with open(quant_file, 'wb') as f:
  f.write(quantized_tflite_model)

with open(float_file, 'wb') as f:
  f.write(float_tflite_model)

print("Float model in Mb:", os.path.getsize(float_file) / float(2**20))
print("Quantized model in Mb:", os.path.getsize(quant_file) / float(2**20))

Conclusão

Neste tutorial, você viu como criar modelos conscientes de quantização com a API de Kit de ferramentas para quantização de modelos do TensorFlow e, em seguida, como quantizar modelos para o back-end do TFLite.

Você viu o benefício da redução do tamanho do modelo em 4x para o MNIST, com uma diferença mínima na exatidão. Para conferir os benefícios da latência em dispositivos móveis, teste os exemplos no repositório de apps do TFLite.

Recomendamos que você teste essa nova capacidade que pode ser especialmente importante para a implantação em ambientes com recursos limitados.