Cuantificación de float16 posterior al entrenamiento

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

Visión general

TensorFlow Lite ahora soporta la conversión de pesos a valores de coma flotante de 16 bits durante la conversión de modelo de TensorFlow a formato tampón plana de TensorFlow Lite. Esto da como resultado una reducción de 2x en el tamaño del modelo. Algunos hardware, como las GPU, pueden computar de forma nativa en esta aritmética de precisión reducida, logrando una aceleración con respecto a la ejecución tradicional de punto flotante. El delegado de la GPU de Tensorflow Lite se puede configurar para que se ejecute de esta manera. Sin embargo, un modelo convertido a pesos float16 aún se puede ejecutar en la CPU sin modificaciones adicionales: los pesos float16 se muestrean a float32 antes de la primera inferencia. Esto permite una reducción significativa en el tamaño del modelo a cambio de un impacto mínimo en la latencia y la precisión.

En este instructivo, entrenará un modelo MNIST desde cero, verificará su precisión en TensorFlow y luego convertirá el modelo en un búfer plano de Tensorflow Lite con cuantificación float16. Finalmente, verifique la precisión del modelo convertido y compárelo con el modelo float32 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
tf.float16
tf.float16

Entrene y exporte el modelo

# 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)
)
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step
11501568/11490434 [==============================] - 0s 0us/step
2021-09-12 11:11:37.167262: 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:11:37.167422: 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.2643 - accuracy: 0.9280 - val_loss: 0.1068 - val_accuracy: 0.9681
<keras.callbacks.History at 0x7f173f943f90>

Por ejemplo, entrenó el modelo solo durante una época, por lo que 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:11:48.574339: 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/tmpqgyupoty/assets
2021-09-12 11:11:49.006295: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:351] Ignored output_format.
2021-09-12 11:11:49.006342: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:354] Ignored drop_control_dependency.

Escribirlo a un .tflite archivo:

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 lugar el modelo de float16 en la exportación, establecido por primera vez la optimizations la bandera de optimizaciones uso por defecto. Luego, especifique que float16 es el tipo admitido en la plataforma de destino:

converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]

Finalmente, convierta el modelo como de costumbre. Tenga en cuenta que, de forma predeterminada, el modelo convertido seguirá utilizando entradas y salidas flotantes para la conveniencia de la invocación.

tflite_fp16_model = converter.convert()
tflite_model_fp16_file = tflite_models_dir/"mnist_model_quant_f16.tflite"
tflite_model_fp16_file.write_bytes(tflite_fp16_model)
INFO:tensorflow:Assets written to: /tmp/tmpw53m2y4_/assets
INFO:tensorflow:Assets written to: /tmp/tmpw53m2y4_/assets
2021-09-12 11:11:49.600290: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:351] Ignored output_format.
2021-09-12 11:11:49.600336: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:354] Ignored drop_control_dependency.
44432

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

ls -lh {tflite_models_dir}
total 128K
-rw-rw-r-- 1 kbuilder kbuilder 83K Sep 12 11:11 mnist_model.tflite
-rw-rw-r-- 1 kbuilder kbuilder 44K Sep 12 11:11 mnist_model_quant_f16.tflite

Ejecute los modelos de TensorFlow Lite

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

Cargue el modelo en los intérpretes

interpreter = tf.lite.Interpreter(model_path=str(tflite_model_file))
interpreter.allocate_tensors()
interpreter_fp16 = tf.lite.Interpreter(model_path=str(tflite_model_fp16_file))
interpreter_fp16.allocate_tensors()

Pruebe los modelos 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

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

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

interpreter_fp16.set_tensor(input_index, test_image)
interpreter_fp16.invoke()
predictions = interpreter_fp16.get_tensor(output_index)
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.9681

Repita la evaluación en el modelo cuantificado float16 para obtener:

# NOTE: Colab runs on server CPUs. At the time of writing this, TensorFlow Lite
# doesn't have super optimized server CPU kernels. For this reason this may be
# slower than the above float interpreter. But for mobile CPUs, considerable
# speedup can be observed.
print(evaluate_model(interpreter_fp16))
0.9681

En este ejemplo, ha cuantificado un modelo para que flote16 sin diferencia en la precisión.

También es posible evaluar el modelo cuantificado fp16 en la GPU. Para llevar a cabo toda la aritmética con los valores de precisión reducidos, asegúrese de crear las TfLiteGPUDelegateOptions struct en su aplicación y conjunto precision_loss_allowed a 1 , de esta manera:

//Prepare GPU delegate.
const TfLiteGpuDelegateOptions options = {
  .metadata = NULL,
  .compile_options = {
    .precision_loss_allowed = 1,  // FP16
    .preferred_gl_object_type = TFLITE_GL_OBJECT_TYPE_FASTEST,
    .dynamic_batch_enabled = 0,   // Not fully functional yet
  },
};

Documentación detallada sobre el delegado TFLite GPU y la forma de utilizarlo en su aplicación se puede encontrar aquí