TensorFlow Lite转换器

TensorFlow Lite转换器

TensorFlow Lite 转换器是将TensorFlow模型编译、生成可供移动设备使用的TensorFlow Lite FlatBuffer 文件（.tflite）。该转换器支持 SavedModel 目录、tf.keras 模型和具体函数。
一张图看清模型转换的作用：
转换方法有两种：Python API 、命令行方式。

一、Python API

在 TensorFlow 2.0 中，用来将原始的 TensorFlow 模型格式转换为 TensorFlow Lite 的 Python API 是 tf.lite.TFLiteConverter。在 TFLiteConverter 中有以下的类方法（classmethod）：

• TFLiteConverter.from_saved_model()：用来转换 SavedModel 格式模型。
• TFLiteConverter.from_keras_model()：用来转换 tf.keras 模型。
• TFLiteConverter.from_concrete_functions()：用来转换 concrete functions。

注意: 在 TensorFlow Lite 2.0 中有一个不同版本的 TFLiteConverter API， 该 API 只包含了 from_concrete_function。 本文中用到的的新版本 API 可以通过 pip 安装 tf-nightly-2.0-preview。

示例：

转换 SavedModel 格式模型

以下示例展示了如何将一个 SavedModel 转换为 TensorFlow Lite 中的 FlatBuffer格式。

import tensorflow as tf

# 建立一个简单的模型。
root = tf.train.Checkpoint()
root.v1 = tf.Variable(3.)
root.v2 = tf.Variable(2.)
root.f = tf.function(lambda x: root.v1 * root.v2 * x)

# 保存模型。
export_dir = "/tmp/test_saved_model"
input_data = tf.constant(1., shape=[1, 1])
to_save = root.f.get_concrete_function(input_data)
tf.saved_model.save(root, export_dir, to_save)

# 转换模型。
converter = tf.lite.TFLiteConverter.from_saved_model(export_dir)
tflite_model = converter.convert()

此 API 不支持指定输入向量的维度。 如果您的模型需要指定输入向量的维度，请使用 from_concrete_functions 来完成。 示例：

model = tf.saved_model.load(export_dir)
concrete_func = model.signatures[
  tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
concrete_func.inputs[0].set_shape([1, 256, 256, 3])
converter = TFLiteConverter.from_concrete_functions([concrete_func])

转换 Keras 模型

以下示例展示了如何将一个 tf.keras 模型 转换为 TensorFlow Lite 中的 FlatBuffer 格式。

import tensorflow as tf

# 创建一个简单的 Keras 模型。
x = [-1, 0, 1, 2, 3, 4]
y = [-3, -1, 1, 3, 5, 7]

model = tf.keras.models.Sequential(
    [tf.keras.layers.Dense(units=1, input_shape=[1])])
model.compile(optimizer='sgd', loss='mean_squared_error')
model.fit(x, y, epochs=50)

# 转换模型。
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

转换 concrete function

以下示例展示了如何将 TensorFlow 中的 concrete function 转换为TensorFlow Lite 中的 FlatBuffer 格式。

import tensorflow as tf

# 建立一个模型。
root = tf.train.Checkpoint()
root.v1 = tf.Variable(3.)
root.v2 = tf.Variable(2.)
root.f = tf.function(lambda x: root.v1 * root.v2 * x)

# 生成 concrete function。
input_data = tf.constant(1., shape=[1, 1])
concrete_func = root.f.get_concrete_function(input_data)

# 转换模型。
#
# `from_concrete_function` 的传入参数被设计为一个一个 concrete function 的列表，然而
# 现阶段仅支持每次调用时仅接受一个concrete function。
# 同时转换多个concrete function的功能正在开发中。
converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
tflite_model = converter.convert()

端到端 MobileNet 转换

以下示例展示了如何将将一个提前训练好的 tf.keras MobileNet 模型转换为 TensorFlow Lite 支持的类型并运行推断 （inference）。 随机数据分别在 TensorFlow 和 TensorFlow Lite 模型中运行的结果将被比较。如果是从文件加载模型，请使用 model_path 来代替 model_content。

import numpy as np
import tensorflow as tf

# 加载 MobileNet tf.keras 模型。
model = tf.keras.applications.MobileNetV2(
    weights="imagenet", input_shape=(224, 224, 3))

# 转换模型。
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# 加载 TFLite 模型并分配张量（tensor）。
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

# 获取输入和输出张量。
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 使用随机数据作为输入测试 TensorFlow Lite 模型。
input_shape = input_details[0]['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()` 获取指向张量的指针。
tflite_results = interpreter.get_tensor(output_details[0]['index'])

# 使用随机数据作为输入测试 TensorFlow 模型。
tf_results = model(tf.constant(input_data))

# 对比结果。
for tf_result, tflite_result in zip(tf_results, tflite_results):
  np.testing.assert_almost_equal(tf_result, tflite_result, decimal=5)

二、命令行

下列命令参数用于输入和输出文件。

• –output_file. 类型: string. 指定输出文件的绝对路径。
• –saved_model_dir. 类型: string. 指定含有 TensorFlow 1.x 或者 2.0 使用 SavedModel 生成文件的绝对路径目录。
• –keras_model_file. Type: string. 指定含有 TensorFlow 1.x 或者 2.0 使用 tf.keras model 生成 HDF5 文件的绝对路径目录。

tflite_convert \
  --saved_model_dir=/tmp/mobilenet_saved_model \
  --output_file=/tmp/mobilenet.tflite