Cuantización de rango dinámico posterior al entrenamiento

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar cuaderno Ver modelo TF Hub

Visión general

TensorFlow Lite ahora soporta la conversión de pesos a la precisión de 8 bits como parte de la conversión del modelo graphdefs tensorflow al formato del búfer plana de TensorFlow Lite. La cuantificación de rango dinámico logra una reducción 4x en el tamaño del modelo. Además, TFLite admite la cuantificación y descuantificación sobre la marcha de activaciones para permitir:

  1. Uso de núcleos cuantificados para una implementación más rápida cuando esté disponible.
  2. Mezcla de núcleos de punto flotante con núcleos cuantificados para diferentes partes del gráfico.

Las activaciones siempre se almacenan en coma flotante. Para operaciones que admiten kernels cuantificados, las activaciones se cuantifican a 8 bits de precisión dinámicamente antes del procesamiento y se descuantifican para flotar con precisión después del procesamiento. Dependiendo del modelo que se convierta, esto puede acelerar el cálculo de punto flotante puro.

En contraste con la cuantificación de formación conscientes , los pesos son formación post cuantificado y las activaciones se cuantifican de forma dinámica en la inferencia en este método. Por lo tanto, los pesos del modelo no se vuelven a entrenar para compensar los errores inducidos por cuantificación. Es importante verificar la precisión del modelo cuantificado para asegurarse de que la degradación sea aceptable.

Este tutorial entrena un modelo MNIST desde cero, verifica su precisión en TensorFlow y luego convierte el modelo en un búfer plano de Tensorflow Lite con cuantificación de rango dinámico. Finalmente, verifica la precisión del modelo convertido y lo compara con el modelo flotante original.

Construye un modelo MNIST

Configuración

import logging
logging.getLogger("tensorflow").setLevel(logging.DEBUG)

import tensorflow as tf
from tensorflow import keras
import numpy as np
import pathlib

Entrena un modelo de TensorFlow

# 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=tf.nn.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=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
model.fit(
  train_images,
  train_labels,
  epochs=1,
  validation_data=(test_images, test_labels)
)
2021-09-12 11:15:11.173006: E tensorflow/stream_executor/cuda/cuda_driver.cc:271] failed call to cuInit: CUDA_ERROR_SYSTEM_DRIVER_MISMATCH: system has unsupported display driver / cuda driver combination
2021-09-12 11:15:11.173171: E tensorflow/stream_executor/cuda/cuda_diagnostics.cc:313] kernel version 470.57.2 does not match DSO version 470.63.1 -- cannot find working devices in this configuration
1875/1875 [==============================] - 11s 5ms/step - loss: 0.2621 - accuracy: 0.9256 - val_loss: 0.1208 - val_accuracy: 0.9627
<keras.callbacks.History at 0x7fc628be2550>

Por ejemplo, dado que entrenó el modelo solo durante una época, solo entrena con una precisión de ~ 96%.

Convertir a un modelo de TensorFlow Lite

Utilizando el pitón TFLiteConverter , ahora se puede convertir el modelo entrenado en un modelo TensorFlow Lite.

Ahora carga el modelo usando el TFLiteConverter :

converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
2021-09-12 11:15:22.477956: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
INFO:tensorflow:Assets written to: /tmp/tmphu4bq30h/assets
2021-09-12 11:15:22.930564: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:351] Ignored output_format.
2021-09-12 11:15:22.930614: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:354] Ignored drop_control_dependency.

Escríbalo en un archivo tflite:

tflite_models_dir = pathlib.Path("/tmp/mnist_tflite_models/")
tflite_models_dir.mkdir(exist_ok=True, parents=True)
tflite_model_file = tflite_models_dir/"mnist_model.tflite"
tflite_model_file.write_bytes(tflite_model)
84500

Para cuantificar el modelo de exportación, establecer la optimizations la bandera para optimizar el tamaño:

converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()
tflite_model_quant_file = tflite_models_dir/"mnist_model_quant.tflite"
tflite_model_quant_file.write_bytes(tflite_quant_model)
INFO:tensorflow:Assets written to: /tmp/tmp2j3lul1o/assets
INFO:tensorflow:Assets written to: /tmp/tmp2j3lul1o/assets
2021-09-12 11:15:23.566446: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:351] Ignored output_format.
2021-09-12 11:15:23.566494: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:354] Ignored drop_control_dependency.
23904

Nótese cómo el archivo resultante, es de aproximadamente 1/4 del tamaño.

ls -lh {tflite_models_dir}
total 180K
-rw-rw-r-- 1 kbuilder kbuilder 83K Sep 12 11:15 mnist_model.tflite
-rw-rw-r-- 1 kbuilder kbuilder 24K Sep 12 11:15 mnist_model_quant.tflite
-rw-rw-r-- 1 kbuilder kbuilder 25K Sep 12 11:14 mnist_model_quant_16x8.tflite
-rw-rw-r-- 1 kbuilder kbuilder 44K Sep 12 11:11 mnist_model_quant_f16.tflite

Ejecute los modelos TFLite

Ejecute el modelo de TensorFlow Lite con el intérprete de Python TensorFlow Lite.

Cargue el modelo en un intérprete

interpreter = tf.lite.Interpreter(model_path=str(tflite_model_file))
interpreter.allocate_tensors()
interpreter_quant = tf.lite.Interpreter(model_path=str(tflite_model_quant_file))
interpreter_quant.allocate_tensors()

Prueba el modelo en una imagen

test_image = np.expand_dims(test_images[0], axis=0).astype(np.float32)

input_index = interpreter.get_input_details()[0]["index"]
output_index = interpreter.get_output_details()[0]["index"]

interpreter.set_tensor(input_index, test_image)
interpreter.invoke()
predictions = interpreter.get_tensor(output_index)
import matplotlib.pylab as plt

plt.imshow(test_images[0])
template = "True:{true}, predicted:{predict}"
_ = plt.title(template.format(true= str(test_labels[0]),
                              predict=str(np.argmax(predictions[0]))))
plt.grid(False)

png

Evaluar los modelos

# A helper function to evaluate the TF Lite model using "test" dataset.
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 test_image in test_images:
    # 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)

  # Compare prediction results with ground truth labels to calculate accuracy.
  accurate_count = 0
  for index in range(len(prediction_digits)):
    if prediction_digits[index] == test_labels[index]:
      accurate_count += 1
  accuracy = accurate_count * 1.0 / len(prediction_digits)

  return accuracy
print(evaluate_model(interpreter))
0.9627

Repita la evaluación en el modelo cuantificado de rango dinámico para obtener:

print(evaluate_model(interpreter_quant))
0.9631

En este ejemplo, el modelo comprimido no tiene diferencia en la precisión.

Optimización de un modelo existente

Los resnets con capas de preactivación (Resnet-v2) se utilizan ampliamente para aplicaciones de visión. Pre-entrenado gráfico congelada para resnet-v2-101 está disponible en Tensorflow Hub .

Puede convertir el gráfico congelado en un búfer plano TensorFLow Lite con cuantificación mediante:

import tensorflow_hub as hub

resnet_v2_101 = tf.keras.Sequential([
  keras.layers.InputLayer(input_shape=(224, 224, 3)),
  hub.KerasLayer("https://tfhub.dev/google/imagenet/resnet_v2_101/classification/4")
])

converter = tf.lite.TFLiteConverter.from_keras_model(resnet_v2_101)
# Convert to TF Lite without quantization
resnet_tflite_file = tflite_models_dir/"resnet_v2_101.tflite"
resnet_tflite_file.write_bytes(converter.convert())
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
INFO:tensorflow:Assets written to: /tmp/tmp0ylf7auc/assets
INFO:tensorflow:Assets written to: /tmp/tmp0ylf7auc/assets
2021-09-12 11:15:50.335557: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:351] Ignored output_format.
2021-09-12 11:15:50.335615: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:354] Ignored drop_control_dependency.
178509352
# Convert to TF Lite with quantization
converter.optimizations = [tf.lite.Optimize.DEFAULT]
resnet_quantized_tflite_file = tflite_models_dir/"resnet_v2_101_quantized.tflite"
resnet_quantized_tflite_file.write_bytes(converter.convert())
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
INFO:tensorflow:Assets written to: /tmp/tmpb5zp9x41/assets
INFO:tensorflow:Assets written to: /tmp/tmpb5zp9x41/assets
2021-09-12 11:16:05.487162: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:351] Ignored output_format.
2021-09-12 11:16:05.487213: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:354] Ignored drop_control_dependency.
46256864
ls -lh {tflite_models_dir}/*.tflite
-rw-rw-r-- 1 kbuilder kbuilder  83K Sep 12 11:15 /tmp/mnist_tflite_models/mnist_model.tflite
-rw-rw-r-- 1 kbuilder kbuilder  24K Sep 12 11:15 /tmp/mnist_tflite_models/mnist_model_quant.tflite
-rw-rw-r-- 1 kbuilder kbuilder  25K Sep 12 11:14 /tmp/mnist_tflite_models/mnist_model_quant_16x8.tflite
-rw-rw-r-- 1 kbuilder kbuilder  44K Sep 12 11:11 /tmp/mnist_tflite_models/mnist_model_quant_f16.tflite
-rw-rw-r-- 1 kbuilder kbuilder 171M Sep 12 11:15 /tmp/mnist_tflite_models/resnet_v2_101.tflite
-rw-rw-r-- 1 kbuilder kbuilder  45M Sep 12 11:16 /tmp/mnist_tflite_models/resnet_v2_101_quantized.tflite

El tamaño del modelo se reduce de 171 MB a 43 MB. La precisión de este modelo en IMAGEnet puede ser evaluada utilizando los scripts proporcionados para medición de precisión TFLite .

La precisión top-1 del modelo optimizado es 76,8, lo mismo que el modelo de punto flotante.