使用 TensorFlow Lite 部署模型

环境配置

线上环境使用的 Tensorflow 版本为 2.1.0。

TensorFlow Lite 模型转换

在 TensorFlow 2 中，用于将 TensorFlow 模型转换为 TensorFlow Lite 的 API 为 tf.lite.TFLiteConverter，其中包含三种方法：
from_keras_model，用于转换 Keras 模型。
from_saved_model，用于转换 SavedModel 模型。
from_concrete_functions，用于转换 TensorFlow 的 Function。

在本节实验中主要使用 from_keras_model 方法。
首先下载预训练的 MobileNetV2 模型，并放到本地的 /root/.keras/models 目录下。

# 下载预训练模型到线上环境指定目录
!wget -nc "https://labfile.oss.aliyuncs.com/courses/1435/mobilenet_v2_weights_tf_dim_ordering_tf_kernels_1.0_224.h5" -P "/root/.keras/models"
!wget -nc "https://labfile.oss.aliyuncs.com/courses/1435/imagenet_class_index.json" -P "/root/.keras/models"

接下来就可以将预训练模型转为 TensorFlow Lite 的格式了。该过程持续时间较长，请耐心等待。

import tensorflow as tf

# 导入 ImageNet 预训练的 MobileNetV2
model = tf.keras.applications.mobilenet_v2.MobileNetV2(weights="imagenet")

# 使用 TFLiteConverter 加载模型并转换
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# 保存为 converted_model.tflite 文件
open("converted_model.tflite", "wb").write(tflite_model)

输出数值代表了模型文件的大小。

基本用法

和之前的实验一样，先准备图片数据。

from skimage import data
import cv2
import numpy as np
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input

# 读取图片并预处理
image = data.chelsea()
image = cv2.resize(image, (224, 224))
image = np.expand_dims(image, 0)
x = preprocess_input(image)
x.shape

用 tf.lite.Interpreter 方法来读取模型文件，并进行预测。

from tensorflow.keras.applications.mobilenet_v2 import decode_predictions


def pred_fn(filename):
    # 从文件中加载 TFLite 模型并分配张量
    interpreter = tf.lite.Interpreter(model_path=filename)
    interpreter.allocate_tensors()
    # 获取输入和输出张量
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    # 张量填充
    interpreter.set_tensor(input_details[0]['index'], x)
    # 运行推理
    interpreter.invoke()
    # 获取张量值
    results = interpreter.get_tensor(output_details[0]['index'])
    # 获取对应标签
    pred = decode_predictions(results)
    return pred

pred_fn('converted_model.tflite')

上述代码中 TensorFlow Lite 的工作主要分为初始化和推理。
初始化用于减少推理时的操作，保证模型推理高效运行，其中包括模型文件加载，获取输入输出张量。
推理进行的操作包括输入张量填充，interpreter.invoke，输出张量获取。

TensorFlow Lite 优化与加速

使用 TensorFlow Lite 解释器进行推断。首先，将从 TensorFlow 中导入 tflite 改为从 tflite_runtime 中导入，然后直接从 tflite 中获取 Interpreter 并完成推理过程。

# 直接从 tf.lite 中获取 Interpreter
interpreter = tf.lite.Interpreter(model_path='converted_model.tflite')
interpreter.allocate_tensors()
# 获取输入和输出张量
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
# 张量填充
interpreter.set_tensor(input_details[0]['index'], x)
# 运行推理
interpreter.invoke()
# 获取张量值
results = interpreter.get_tensor(output_details[0]['index'])
# 获取对应标签
pred = decode_predictions(results)
pred

模型量化
移动设备一般只有比较低的计算能力，而且也会受限于内存和电量消耗。也有必要降低将模型传输给设备的通信量及降低网络连接需求。因此，将模型量化使其模型尺寸更小、推断更快、耗电更低是非常有必要的。训练时通常使用 Float32 格式的数据，训练后量化可以减小模型大小，同时还可以改善 CPU 和硬件加速器的延迟。通常，训练后量化有以下三种：

image.png

# 使用 TFLiteConverter 加载模型
weight_converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 使用 optimizations 对输出大小进行优化
weight_converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
# 转换模型并保存
tflite_weight_model = weight_converter.convert()
open("weight_converted_model.tflite", "wb").write(tflite_weight_model)

转换完成后，就可以使用上面导入的图片进行测试了。

pred_fn("weight_converted_model.tflite")

整数量化
整数量化将模型所有的值（权重和激活）转换为 8 位整数，在这里为了创建具有准确的动态激活范围的量化值，还需要提供一个代表性的数据集。

# 生成代表性的数据集，在实际中可能需要部分或者全部的训练数据
dataset = tf.data.Dataset.from_tensor_slices((x)).batch(1)
def representative_dataset_gen():
  for input_value in dataset.take(1):
    yield [input_value]

再次使用量化转换模型。

# 使用 TFLiteConverter 加载模型
integer_converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 使用 optimizations 对输出大小进行优化
integer_converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 指定代表性的数据集
integer_converter.representative_dataset = representative_dataset_gen
# 转换并保存
tflite_integer_model = integer_converter.convert()
open("integer_converted_model.tflite", "wb").write(tflite_integer_model)

# 使用图片测试输出结果
pred_fn('integer_converted_model.tflite')

Float16 量化
Float16 量化将权重转换为 16bit 浮点值。这样可将模型尺寸减小 2 倍，一些硬件（例如 GPU）可以使用这种降低精度的算法进行计算，从而与传统浮点运算相比具有更快的执行速度。

# 使用 TFLiteConverter 加载模型
float16_converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 使用 optimizations 对输出大小进行优化，同时指定 float16 为目标设备上的支持类型
float16_converter.optimizations = [tf.lite.Optimize.DEFAULT]
float16_converter.target_spec.supported_types = [tf.float16]
# 转换并保存
tflite_float16_model = float16_converter.convert()
open("float16_converted_model.tflite", "wb").write(tflite_float16_model)

# 使用图片测试输出结果
pred_fn("float16_converted_model.tflite")

最后，我们可以查看对比一下所有文件的大小。

!ls -Slh *.tflite /root/.keras/models/*.h5

由此可见，量化对于模型压缩的效果是十分明显的。