TensorFlow Lite 转换器是将TensorFlow模型编译、生成可供移动设备使用的TensorFlow Lite FlatBuffer 文件(.tflite)。该转换器支持 SavedModel 目录、tf.keras 模型和具体函数。
一张图看清模型转换的作用:
转换方法有两种:Python API 、命令行方式。
在 TensorFlow 2.0 中,用来将原始的 TensorFlow 模型格式转换为 TensorFlow Lite 的 Python API 是 tf.lite.TFLiteConverter。在 TFLiteConverter 中有以下的类方法(classmethod):
注意: 在 TensorFlow Lite 2.0 中有一个不同版本的 TFLiteConverter API, 该 API 只包含了 from_concrete_function。 本文中用到的的新版本 API 可以通过 pip 安装 tf-nightly-2.0-preview。
以下示例展示了如何将一个 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])
以下示例展示了如何将一个 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()
以下示例展示了如何将 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()
以下示例展示了如何将将一个提前训练好的 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)
下列命令参数用于输入和输出文件。
tflite_convert \
--saved_model_dir=/tmp/mobilenet_saved_model \
--output_file=/tmp/mobilenet.tflite