訓練後整數量化

在 TensorFlow.org 上查看

在 Google Colab 中執行

在 GitHub 上查看原始碼

下載筆記本

總覽

整數量化是一種最佳化策略，可將 32 位元浮點數 (例如權重和啟動輸出) 轉換為最接近的 8 位元定點數。這樣可以縮減模型大小並提高推論速度，這對於 微控制器等低功耗裝置而言非常寶貴。整數專用加速器 (例如 Edge TPU) 也需要這種資料格式。

在本教學課程中，您將從頭開始訓練 MNIST 模型、將其轉換為 Tensorflow Lite 檔案，並使用訓練後量化進行量化。最後，您將檢查已轉換模型的準確度，並將其與原始浮點模型進行比較。

實際上，您可以選擇模型量化的程度。在本教學課程中，您將執行「完整整數量化」，將所有權重和啟動輸出轉換為 8 位元整數資料，而其他策略可能會保留一些浮點資料。

如要進一步瞭解各種量化策略，請參閱TensorFlow Lite 模型最佳化。

設定

為了量化輸入和輸出張量，我們需要使用 TensorFlow 2.3 中新增的 API

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

import tensorflow as tf
import numpy as np
print("TensorFlow version: ", tf.__version__)

產生 TensorFlow 模型

我們將建構一個簡單的模型，以分類來自 MNIST 資料集的數字。

這個訓練不會花費太長時間，因為您只針對 5 個週期訓練模型，這大約可訓練到 ~98% 的準確度。

# 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.astype(np.float32) / 255.0
test_images = test_images.astype(np.float32) / 255.0

# Define the model architecture
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='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,
  epochs=5,
  validation_data=(test_images, test_labels)
)

轉換為 TensorFlow Lite 模型

現在，您可以使用 TensorFlow Lite Converter 將已訓練模型轉換為 TensorFlow Lite 格式，並套用不同程度的量化。

請注意，某些版本的量化會將部分資料保留為浮點格式。因此，以下章節將逐步展示每個選項，量化程度會逐漸增加，直到我們獲得完全是 int8 或 uint8 資料的模型。(請注意，我們在每個章節中重複了一些程式碼，以便您查看每個選項的所有量化步驟。)

首先，以下是未經量化的已轉換模型

converter = tf.lite.TFLiteConverter.from_keras_model(model)

tflite_model = converter.convert()

現在它是一個 TensorFlow Lite 模型，但所有參數資料仍然使用 32 位元浮點值。

使用動態範圍量化進行轉換

現在，讓我們啟用預設 optimizations 旗標來量化所有固定參數 (例如權重)

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

tflite_model_quant = converter.convert()

模型現在變得更小，並且權重已量化，但其他可變資料仍為浮點格式。

使用浮點回退量化進行轉換

為了量化可變資料 (例如模型輸入/輸出和層之間的資料)，您需要提供 RepresentativeDataset。這是一個產生器函式，可提供一組夠大的輸入資料來表示典型值。它可讓轉換器估計所有可變資料的動態範圍。(與訓練或評估資料集相比，資料集不需要是唯一的。) 為了支援多個輸入，每個代表性資料點都是一個清單，清單中的元素會根據其索引饋送到模型。

def representative_data_gen():
  for input_value in tf.data.Dataset.from_tensor_slices(train_images).batch(1).take(100):
    # Model has only one input so each data point has one element.
    yield [input_value]

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen

tflite_model_quant = converter.convert()

現在，所有權重和可變資料都已量化，而且與原始 TensorFlow Lite 模型相比，模型明顯更小。

但是，為了維持與傳統上使用浮點模型輸入和輸出張量的應用程式的相容性，TensorFlow Lite Converter 會將模型輸入和輸出張量保留為浮點

interpreter = tf.lite.Interpreter(model_content=tflite_model_quant)
input_type = interpreter.get_input_details()[0]['dtype']
print('input: ', input_type)
output_type = interpreter.get_output_details()[0]['dtype']
print('output: ', output_type)

這通常有利於相容性，但無法與僅執行以整數為基礎的運算的裝置 (例如 Edge TPU) 相容。

此外，如果 TensorFlow Lite 未包含該運算的量化實作，則上述程序可能會將運算保留為浮點格式。此策略可讓轉換完成，讓您獲得更小且更有效率的模型，但同樣地，它無法與僅限整數的硬體相容。(此 MNIST 模型中的所有運算都有量化實作。)

因此，為了確保端對端僅限整數的模型，您需要再新增幾個參數...

使用僅限整數量化進行轉換

為了量化輸入和輸出張量，並讓轉換器在遇到無法量化的運算時擲回錯誤，請使用一些額外參數再次轉換模型

def representative_data_gen():
  for input_value in tf.data.Dataset.from_tensor_slices(train_images).batch(1).take(100):
    yield [input_value]

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen
# Ensure that if any ops can't be quantized, the converter throws an error
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
# Set the input and output tensors to uint8 (APIs added in r2.3)
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8

tflite_model_quant = converter.convert()

內部量化與上述相同，但您可以看到輸入和輸出張量現在是整數格式

interpreter = tf.lite.Interpreter(model_content=tflite_model_quant)
input_type = interpreter.get_input_details()[0]['dtype']
print('input: ', input_type)
output_type = interpreter.get_output_details()[0]['dtype']
print('output: ', output_type)

現在您有一個整數量化模型，該模型針對模型的輸入和輸出張量使用整數資料，因此它與僅限整數的硬體 (例如 Edge TPU) 相容。

將模型儲存為檔案

您需要 .tflite 檔案才能在其他裝置上部署模型。因此，讓我們將已轉換的模型儲存到檔案，然後在我們在下方執行推論時載入它們。

import pathlib

tflite_models_dir = pathlib.Path("/tmp/mnist_tflite_models/")
tflite_models_dir.mkdir(exist_ok=True, parents=True)

# Save the unquantized/float model:
tflite_model_file = tflite_models_dir/"mnist_model.tflite"
tflite_model_file.write_bytes(tflite_model)
# Save the quantized model:
tflite_model_quant_file = tflite_models_dir/"mnist_model_quant.tflite"
tflite_model_quant_file.write_bytes(tflite_model_quant)

執行 TensorFlow Lite 模型

現在，我們將使用 TensorFlow Lite Interpreter 執行推論，以比較模型準確度。

首先，我們需要一個函式，該函式使用給定模型和圖片執行推論，然後傳回預測

# Helper function to run inference on a TFLite model
def run_tflite_model(tflite_file, test_image_indices):
  global test_images

  # Initialize the interpreter
  interpreter = tf.lite.Interpreter(model_path=str(tflite_file))
  interpreter.allocate_tensors()

  input_details = interpreter.get_input_details()[0]
  output_details = interpreter.get_output_details()[0]

  predictions = np.zeros((len(test_image_indices),), dtype=int)
  for i, test_image_index in enumerate(test_image_indices):
    test_image = test_images[test_image_index]

    # Check if the input type is quantized, then rescale input data to uint8
    if input_details['dtype'] == np.uint8:
      input_scale, input_zero_point = input_details["quantization"]
      test_image = test_image / input_scale + input_zero_point

    test_image = np.expand_dims(test_image, axis=0).astype(input_details["dtype"])
    interpreter.set_tensor(input_details["index"], test_image)
    interpreter.invoke()
    output = interpreter.get_tensor(output_details["index"])[0]

    predictions[i] = output.argmax()

  return predictions

在單一圖片上測試模型

現在，我們將比較浮點模型和量化模型的效能

• tflite_model_file 是具有浮點資料的原始 TensorFlow Lite 模型。
• tflite_model_quant_file 是我們使用僅限整數量化轉換的最後一個模型 (它針對輸入和輸出使用 uint8 資料)。

讓我們建立另一個函式來列印我們的預測

import matplotlib.pylab as plt

# Change this to test a different image
test_image_index = 1

## Helper function to test the models on one image
def test_model(tflite_file, test_image_index, model_type):
  global test_labels

  predictions = run_tflite_model(tflite_file, [test_image_index])

  plt.imshow(test_images[test_image_index])
  template = model_type + " Model \n True:{true}, Predicted:{predict}"
  _ = plt.title(template.format(true= str(test_labels[test_image_index]), predict=str(predictions[0])))
  plt.grid(False)

現在測試浮點模型

test_model(tflite_model_file, test_image_index, model_type="Float")

並測試量化模型

test_model(tflite_model_quant_file, test_image_index, model_type="Quantized")

評估所有圖片上的模型

現在，讓我們使用我們在本教學課程開始時載入的所有測試圖片來執行這兩個模型

# Helper function to evaluate a TFLite model on all images
def evaluate_model(tflite_file, model_type):
  global test_images
  global test_labels

  test_image_indices = range(test_images.shape[0])
  predictions = run_tflite_model(tflite_file, test_image_indices)

  accuracy = (np.sum(test_labels== predictions) * 100) / len(test_images)

  print('%s model accuracy is %.4f%% (Number of test samples=%d)' % (
      model_type, accuracy, len(test_images)))

評估浮點模型

evaluate_model(tflite_model_file, model_type="Float")

評估量化模型

evaluate_model(tflite_model_quant_file, model_type="Quantized")

因此，與浮點模型相比，您現在擁有一個整數量化模型，其準確度幾乎沒有差異。

如要進一步瞭解其他量化策略，請參閱TensorFlow Lite 模型最佳化。