84TensorFlow 2 模型部署方法实践--使用 TensorFlow Lite 部署模型

使用 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

量化原理
目前深度学习中大部分都是用 32bit Float 类型进行计算的，bit 位数的多少直接限制了数据类型能够表达的数据范围，比如 Float32 的数据是由 1bit表示符号，8bit 表示整数部，23bit 表示分数部组成。
用更低位的数值类型意味着更小的数据表示范围和更稀疏的数值，量化的时候就会造成数值精度损失。比如要把 Float 数值量化到 Int 类型，那么首先小数部分会损失，而那些超过 Int 类型能够表达的范围的值也会被压缩到 Int 能够表达的最大或最小值。
但是因为芯片的设计原因，低 bit 位的数值计算一般会比高 bit 的要快。同时从 RAM 中读取低 bit 位的数据的时间会缩短，也能降低运算能耗。那么，为什么量化能够在深度学习中应用呢？主要原因是因为训练好的神经网络应该具有一定的抗噪抗扰动能力，且大部分权重都落在一个很小的区间内。
接下来我们将讲解如何使用 TensorFlow Lite 进行模型量化，其内部原理可以参考官方发布的白皮书 Quantizing Deep Convolutional Networks For Efficient Inference。
权重量化
权重量化（也称为混合量化）将权重转换为 8 位精度，权重量化可以使模型尺寸减少 4 倍。

# 使用 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

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