Tensorflow Lite从入门到精通

时间：12-08来源：作者：点击数：

TensorFlow Lite 是 TensorFlow 在移动和 IoT 等边缘设备端的解决方案，提供了 Java、Python 和 C++ API 库，可以运行在 Android、iOS 和 Raspberry Pi 等设备上。目前 TFLite 只提供了推理功能，在服务器端进行训练后，经过如下简单处理即可部署到边缘设备上。

个人使用总结：

如果我们只使用Tensorflow的高级API搭建模型，那么将TF转TF Lite再转TF lite micro的过程会相对顺利。但是如果我们的模型使用了自定义模块，那么转换过程会遇到很多麻烦，Tensorflow对自家高级API的转换提供了很好的支持，但对我们自己写的一些NN 算子支持不佳。
Tensorflow LSTM的流式部署是有难度的，请使用最新的Tensorflow版本，将unrool打开，再尝试

转换模型

我们可以通过以下两种方式将Tensorflow模型转换成 TF Lite：

Python API（推荐，且本文主讲）：它让您可以更轻松地在模型开发流水线中转换模型、应用优化、添加元数据，并且拥有更多功能。
命令行：它仅支持基本模型转换。

我们可以根据保存模型的方式来选择转换成TF lite的方式：

1、tf.lite.TFLiteConverter.from_saved_model()（推荐）：转换 SavedModel

import tensorflow as tf

# 转换模型
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) # SavedModel目录的路径
tflite_model = converter.convert()

# 保存TF Lite模型
with open('model.tflite', 'wb') as f:
  f.write(tflite_model)

2、tf.lite.TFLiteConverter.from_keras_model()：转换 Keras 模型

import tensorflow as tf

# 使用高级API tf.keras.* 创建一个模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(units=1, input_shape=[1]),
    tf.keras.layers.Dense(units=16, activation='relu'),
    tf.keras.layers.Dense(units=1)
])
model.compile(optimizer='sgd', loss='mean_squared_error')  # 编译模型
model.fit(x=[-1, 0, 1], y=[-3, -1, 1], epochs=5)  # 训练模型
# 生成一个 SavedModel 
# tf.saved_model.save(model, "saved_model_keras_dir")

# keras model to TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# 保存模型
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

3、tf.lite.TFLiteConverter.from_concrete_functions()：转换具体函数

import tensorflow as tf


# 使用低级API tf.* 创建一个模型
class Squared(tf.Module):
    @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32)])
    def __call__(self, x):
        return tf.square(x)


model = Squared()
# 运行模型
# result = Squared(5.0) # 25.0
# 生成 SavedModel
# tf.saved_model.save(model, "saved_model_tf_dir")
concrete_func = model.__call__.get_concrete_function()

# TF model to TFLite
# 注意，对于TensorFlow 2.7之前的版本，
# from_concrete_functions API能够在只有第一个参数的情况下工作:
# > converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func], model)
tflite_model = converter.convert()

# 保存TFLite模型
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

转换RNN模型为TFLite

由于 TensorFlow 中 RNN API 的变体很多，我们的转换方式包括两个方面：

为标准 TensorFlow RNN API（如 Keras LSTM）提供原生支持。这是推荐的选项。
提供 进入转换基础架构的接口，用于 插入用户定义的 RNN 实现并转换为 TensorFlow Lite。我们提供了几个有关此类转换的开箱即用的示例，这些示例使用的是 lingvo 的 LSTMCellSimple 和 LayerNormalizedLSTMCellSimple RNN 接口。

Tensorflow官方提供了一个一个非常棒的案例：Keras LSTM fusion Codelab.ipynb

运行推理

TensorFlow Lite 推理通常遵循以下步骤：

加载模型
转换数据
运行推断
解释输出

import numpy as np
import tensorflow as tf

# 加载TFLite模型并分配张量
interpreter = tf.lite.Interpreter(model_path="converted_model.tflite")
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()  # 输入
output_details = interpreter.get_output_details()  # 输出

input_shape = input_details[0]['shape']  # 获取输入的shape
input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32)

interpreter.set_tensor(input_details[0]['index'], input_data)  # 输入给模型
interpreter.invoke()  # 运行
# 函数`get_tensor()`返回张量数据的副本
# 使用`tensor()`来获得一个指向这个张量的指针
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)

定义了Signature的TFLite 推理

# -*- coding:utf-8 -*-
# Author:凌逆战 | Never
# Date: 2022/10/12
"""

"""
import tensorflow as tf


class TestModel(tf.Module):
    def __init__(self):
        super(TestModel, self).__init__()

    @tf.function(input_signature=[tf.TensorSpec(shape=[1, 10], dtype=tf.float32)])
    def add(self, x):
        # 为方便起见，将输出命名为“result”。
        return {'result': x + 4}


saved_model_path = './saved_models'
tflite_path = 'content/test_variable.tflite'

# 保存模型
model = TestModel()
# 您可以省略signatures参数，并且将创建一个名为'serving_default'的默认签名名。
tf.saved_model.save(model, saved_model_path,
                    signatures={'my_signature': model.add.get_concrete_function()})

# TF Model to TFLite
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_path)
tflite_model = converter.convert()
with open(tflite_path, 'wb') as f:
    f.write(tflite_model)

# 加载TFLite模型
interpreter = tf.lite.Interpreter(tflite_path)
# 模型中只定义了一个签名，因此默认情况下它将返回签名
# 如果有多个签名，我们就可以传递这个名字
my_signature = interpreter.get_signature_runner()

# my_signature 可以通过输入作为参数调用
output = my_signature(x=tf.constant([1.0], shape=(1, 10), dtype=tf.float32))
# 'output'是包含推理所有输出的字典,在本例中，我们只有一个输出'result'。
print(output['result'])

优化模型

Tensorflow Lite 和 Tensorflow Model Optimization Toolkit (Tensorflow模型优化工具包)提供了最小优化推理复杂性的工具，可将优化推断的复杂性降至最低。深度神经网络的量化使用了一些技术，这些技术可以降低权重的精确表示，并且降低存储和计算。 TensorFlow Model Optimization Toolkit 目前支持通过量化、剪枝和聚类进行优化。

剪枝：剪枝的工作原理是移除模型中对其预测影响很小的参数。剪枝后的模型在磁盘上的大小相同，并且具有相同的运行时延迟，但可以更高效地压缩。这使剪枝成为缩减模型下载大小的实用技术。

未来，TensorFlow Lite 将降低剪枝后模型的延迟。

聚类：聚类的工作原理是将模型中每一层的权重归入预定数量的聚类中，然后共享属于每个单独聚类的权重的质心值。这就减少了模型中唯一权重值的数量，从而降低了其复杂性。这样一来，就可以更高效地压缩聚类后的模型，从而提供类似于剪枝的部署优势。

开发工作流程

首先，检查托管模型中的模型能否用于您的应用。如果不能，建议从训练后量化工具开始，因为它适用范围广，且无需训练数据。
对于无法达到准确率和延迟目标，或硬件加速器支持很重要的情况，量化感知训练是更好的选择。请参阅 TensorFlow Model Optimization Toolkit 下的其他优化技术。
如果要进一步缩减模型大小，可以在量化模型之前尝试剪枝和/或聚类。

量化

有两种形式的量化：训练后量化和量化感知训练。请从训练后量化开始，因为它更易于使用，尽管量化感知训练在模型准确率方面的表现通常更好。

量化感知训练

等我将代码进行了训练后量化后，再来整这个量化感知训练

训练后量化

量化的工作原理是降低模型参数的精度（默认情况为 32 位浮点数）。这样可以获得较小的模型大小和较快的计算速度。TensorFlow Lite 提供以下量化类型:

技术	数据要求	大小缩减	准确率
训练后 Float16 量化	无数据	高达 50%	轻微的准确率损失
训练后动态范围量化	无数据	高达 75%，速度加快 2-3倍	极小的准确率损失
训练后 int8 量化	无标签的代表性样本	高达 75%，速度加快3+倍	极小的准确率损失
量化感知训练	带标签的训练数据	高达 75%	极小的准确率损失

以下决策树可帮助您仅根据预期的模型大小和准确率来选择要用于模型的量化方案。

动态范围量化

权重(float32) 会在训练后量化为整型(int8)，激活会在推断时动态量化，模型大小缩减至原来的四分之一：

TFLite 支持对激活进行动态量化（激活始终以浮点进行存储。对于支持量化内核的算子，激活会在处理前动态量化为 8 位精度，并在处理后反量化为浮点精度。根据被转换的模型，这可以提供比纯浮点计算更快的速度）以实现以下效果：

在可用时使用量化内核加快实现速度。
将计算图不同部分的浮点内核与量化内核混合。

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
# 启动默认的 optimizations 来量化所有固定参数(权重)
converter.optimizations = [tf.lite.Optimize.DEFAULT] 
tflite_quant_model = converter.convert()

此优化提供的延迟接近全定点推断。但是，输出仍使用浮点进行存储，因此使用动态范围算子的加速小于全定点计算。

我们来吃一个完整的栗子，构建一个MNIST模型，并且对它使用动态范围量化，对比量化前后的精度变化：

# -*- coding:utf-8 -*-
# Author:凌逆战 | Never
# Date: 2022/10/12
"""
动态范围量化
"""
import logging

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

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

# 加载MNIST数据集
mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# 归一化输入图像，使每个像素值在0到1之间。
train_images = train_images / 255.0
test_images = test_images / 255.0

# 定义模型结构
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)
])

# 训练数字分类模型
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))

# TF model to TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

tflite_models_dir = pathlib.Path("./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)  # 84824
# with open('model.tflite', 'wb') as f:
#   f.write(tflite_model)

# 量化模型 ------------------------------------
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)  # 24072

# 将模型加载到解释器中
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()  # 分配张量

# 在单个图像上测试模型
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)
print(predictions)


# 使用“test”数据集评估TF Lite模型
def evaluate_model(interpreter):
    input_index = interpreter.get_input_details()[0]["index"]
    output_index = interpreter.get_output_details()[0]["index"]

    # 对“test”数据集中的每个图像进行预测
    prediction_digits = []
    for test_image in test_images:
        # 预处理：添加batch维度并转换为float32以匹配模型的输入数据格式。
        test_image = np.expand_dims(test_image, axis=0).astype(np.float32)
        interpreter.set_tensor(input_index, test_image)

        interpreter.invoke()  # 运行推理

        # 后处理:去除批尺寸，找到概率最高的数字
        output = interpreter.tensor(output_index)
        digit = np.argmax(output()[0])
        prediction_digits.append(digit)

    # 将预测结果与ground truth 标签进行比较，计算精度。
    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.958
print(evaluate_model(interpreter_quant))  # 0.9579

更多细节请参考：Tensorflow官方文档训练后动态范围量化

全整型量化

整型量化是将32位浮点数转换为8位定点数。这样可以缩减模型大小并加快推理速度，这对低功耗设备（如微控制器）很有价值。仅支持整数的加速器（如 Edge TPU）也需要使用此数据格式。

对于全整数量化，需要校准或估算模型中所有浮点张量的范围，即 (min, max)。与权重和偏差等常量张量不同，模型输入、激活（中间层的输出）和模型输出等变量张量不能校准，除非我们运行几个推断周期。因此，转换器需要一个有代表性的数据集来校准它们。这个数据集可以是训练数据或验证数据的一个小子集（大约 100-500 个样本）。请参阅下面的 representative_dataset() 函数。

从 TensorFlow 2.7 版本开始，您可以通过签名指定代表数据集，示例如下：

def representative_dataset():
  for data in dataset:
    yield {
      "image": data.image,
      "bias": data.bias,
    }

如果给定的 TensorFlow 模型中有多个签名，则可以通过指定签名密钥来指定多个数据集：

def representative_dataset():
  # Feed data set for the "encode" signature.
  for data in encode_signature_dataset:
    yield (
      "encode", {
        "image": data.image,
        "bias": data.bias,
      }
    )

  # Feed data set for the "decode" signature.
  for data in decode_signature_dataset:
    yield (
      "decode", {
        "image": data.image,
        "hint": data.hint,
      },
    )

您可以通过提供输入张量列表来生成代表性数据集：

def representative_dataset():
  for data in tf.data.Dataset.from_tensor_slices((images)).batch(1).take(100):
    yield [tf.dtypes.cast(data, tf.float32)]

从 TensorFlow 2.7 版本开始，我们推荐使用基于签名的方法，而不是基于输入张量列表的方法，因为输入张量排序可以很容易地翻转。

出于测试目的，您可以使用如下所示的虚拟数据集：

def representative_dataset():
    for _ in range(100):
      data = np.random.rand(1, 244, 244, 3)
      yield [data.astype(np.float32)]

1、模型整型量化，输入输出还是浮点

启动默认的 optimizations 只能量化固定参数(如：权重)，但是模型输入/输出和层之间的状态值依然是浮点型的，要量化可变中间值，您需要提供 RepresentativeDataset。这是一个生成器函数，它提供一组足够大的输入数据来代表典型值。converter 可以通过该函数估算所有可变数据的动态范围（相比训练或评估数据集，此数据集不必唯一）为了支持多个输入，每个代表性数据点都是一个列表，并且列表中的元素会根据其索引被馈送到模型。

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 启动默认的 optimizations 来量化所有固定参数(权重)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 要量化可变数据(模型输入/输出 和层之间的中间体)，提供 RepresentativeDataset，来估算所有可变数据的动态范围
converter.representative_dataset = representative_data_gen
tflite_model_quant = converter.convert()

现在，所有权重和可变数据都已量化，并且与原始 TensorFlow Lite 模型相比，该模型要小得多。

但是，为了与传统上使用浮点模型输入和输出张量的应用保持兼容，TensorFlow Lite 转换器将模型的输入和输出张量保留为浮点，这通常对兼容性有利，但它无法兼容执行全整形运算的设备（如 Edge TPU）。

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

2、全整型量化

为了量化输入和输出张量，我们需要使用一些附加参数再次转换模型：

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]    # 先启动默认的optimizations将模型权重进行量化
converter.representative_dataset = representative_data_gen  # 使用代表数据集量化模型中间值
# 如果有任何的 ops不能量化，converter 将抛出错误
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] 
# 将输入和输出tensors设置为int8 类型
converter.inference_input_type = tf.int8   # or tf.uint8
converter.inference_output_type = tf.int8  # or 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)    # <class 'numpy.int8'>
output_type = interpreter.get_output_details()[0]['dtype']
print('output: ', output_type)  # <class 'numpy.int8'>

我们来整一个栗子，您将从头开始训练一个 MNIST 模型、将其转换为 TensorFlow Lite 文件，并使用训练后整型量化。最后，将检查转换后模型的准确率并将其与原始浮点模型进行比较。

# -*- coding:utf-8 -*-
# Author:凌逆战 | Never
# Date: 2022/10/12
"""
整型量化
"""
import logging

# 只显示bug不显示 warning
logging.getLogger("tensorflow").setLevel(logging.DEBUG)

import tensorflow as tf
import numpy as np

# 加载 MNIST 数据集
mnist = tf.keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# Normalize 输入图像，使每个像素值在0到1之间
train_images = train_images.astype(np.float32) / 255.0
test_images = test_images.astype(np.float32) / 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='relu'),
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10)
])

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

# 下面是一个没有量化的转换后模型 -------------------------------
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()


# 下面是全整型量化 -----------------------------------------
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]  # 先启动默认的optimizations将模型权重进行量化
converter.representative_dataset = representative_data_gen  # 使用代表数据集量化模型中间值
# 如果有任何的 ops不能量化，converter 将抛出错误
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
# 将输入和输出tensors设置为int8 类型
converter.inference_input_type = tf.uint8  # or tf.int8
converter.inference_output_type = tf.uint8  # or tf.int8

tflite_model_quant = converter.convert()

import pathlib

tflite_models_dir = pathlib.Path("./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)


# 在TFLite模型上运行推理
def run_tflite_model(tflite_file, test_image_indices):
    global test_images

    # 初始化 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]
        test_label = test_labels[test_image_index]

        # 检查输入类型是否被量化，然后将输入数据缩放到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模型
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")
# Float model accuracy is 97.7500% (Number of test samples=10000)
evaluate_model(tflite_model_quant_file, model_type="Quantized")
# Quantized model accuracy is 97.7000% (Number of test samples=10000)

更多细节请参考：Tensorflow官方文档训练后整型量化

float16量化

现在，TensorFlow Lite 支持在模型从 TensorFlow 转换到 TensorFlow Lite FlatBuffer 格式期间将权重转换为 16 位浮点值。这样可以将模型的大小缩减至原来的二分之一。某些硬件（如 GPU）可以在这种精度降低的算术中以原生方式计算，从而实现比传统浮点执行更快的速度。可以将 Tensorflow Lite GPU 委托配置为以这种方式运行。但是，转换为 float16 权重的模型仍可在 CPU 上运行而无需其他修改：float16 权重会在首次推理前上采样为 float32。这样可以在对延迟和准确率造成最小影响的情况下显著缩减模型大小。

float16 量化的优点如下：

将模型的大小缩减一半（因为所有权重都变成其原始大小的一半）。
实现最小的准确率损失。
支持可直接对 float16 数据进行运算的部分委托（例如 GPU 委托），从而使执行速度比 float32 计算更快。

float16 量化的缺点如下：

它不像对定点数学进行量化那样减少那么多延迟。
默认情况下，float16 量化模型在 CPU 上运行时会将权重值“反量化”为 float32。（请注意，GPU 委托不会执行此反量化，因为它可以对 float16 数据进行运算）

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
tflite_quant_model = converter.convert()

在本教程中，您将从头开始训练一个 MNIST 模型，并在 TensorFlow 中检查其准确率，然后使用 float16 量化将此模型转换为 Tensorflow Lite FlatBuffer 格式。最后，检查转换后模型的准确率，并将其与原始 float32 模型进行比较。

# -*- coding:utf-8 -*-
# Author:凌逆战 | Never
# Date: 2022/10/12
"""

"""
import logging

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

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

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

converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
tflite_models_dir = pathlib.Path("./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)

# float16 量化 ----------------------------------------------
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
tflite_fp16_model = converter.convert()  # 转换为TFLite
tflite_model_fp16_file = tflite_models_dir / "mnist_model_quant_f16.tflite"
tflite_model_fp16_file.write_bytes(tflite_fp16_model)

# 将模型加载到解释器中 ------------------------------
# 没有量化的 TFlite
interpreter = tf.lite.Interpreter(model_path=str(tflite_model_file))
interpreter.allocate_tensors()
# float16 量化的 TFLite
interpreter_fp16 = tf.lite.Interpreter(model_path=str(tflite_model_fp16_file))
interpreter_fp16.allocate_tensors()


# 使用 "test" 数据集评估TF Lite模型
def evaluate_model(interpreter):
    input_index = interpreter.get_input_details()[0]["index"]
    output_index = interpreter.get_output_details()[0]["index"]

    # 对“test”数据集中的每个图像进行预测
    prediction_digits = []
    for test_image in test_images:
        # 预处理: 添加batch维度并转换为float32以匹配模型的输入数据格式
        test_image = np.expand_dims(test_image, axis=0).astype(np.float32)
        interpreter.set_tensor(input_index, test_image)

        interpreter.invoke()  # 运行推理

        # 后处理:去除batch维度 找到概率最高的数字
        output = interpreter.tensor(output_index)
        digit = np.argmax(output()[0])
        prediction_digits.append(digit)

    # 将预测结果与ground truth 标签进行比较，计算精度。
    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.9662
# NOTE: Colab运行在服务器的cpu上。
# 在写这篇文章的时候，TensorFlow Lite还没有超级优化的服务器CPU内核。
# 由于这个原因，它可能比上面的float interpreter要慢
# 但是对于mobile CPUs，可以观察到相当大的加速
print(evaluate_model(interpreter_fp16))  # 0.9662

仅整数：具有 8 位权重的 16 位激活（实验性）

这是一个实验性量化方案。它与“仅整数”方案类似，但会根据激活的范围将其量化为 16 位，权重会被量化为 8 位整数，偏差会被量化为 64 位整数。这被进一步称为 16x8 量化。

这种量化的主要优点是可以显著提高准确率，但只会稍微增加模型大小。

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.representative_dataset = representative_dataset
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8]
tflite_quant_model = converter.convert()

如果模型中的部分算子不支持 16x8 量化，模型仍然可以量化，但不受支持的算子会保留为浮点。要允许此操作，应将以下选项添加到 target_spec 中。

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.representative_dataset = representative_dataset
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
                                       tf.lite.OpsSet.TFLITE_BUILTINS]
tflite_quant_model = converter.convert()

这种量化的缺点是：