（13-3）TensorFlow Lite移动端与嵌入式轻量级开发实战：TensorFlow Lite转换器实战

13.3  TensorFlow Lite转换器实战

通过使用TensorFlow Lite转换器，可以根据输入的TensorFlow模型生成TensorFlow Lite 模型。TensorFlow Lite 模型文件是一种优化的 FlatBuffer 格式，以“.tflite”为文件扩展名。

13.3.1  转换方式

在开发过程中，可以通过以下两种方式使用TensorFlow Lite转换器：

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

在接下来的内容中，将详细讲解这两种转换方式的知识和用法。

1. Python API

在使用Python API方式生成TensorFlow Lite 模型之前，需要先确定已安装TensorFlow的版本，具体方法是请运行如下代码：

print(tf.__version__)

要详细了解 TensorFlow Lite converter API的信息，请运行下面的代码：

print(help(tf.lite.TFLiteConverter))

如果开发者已经安装了TensorFlow，则可以使用tf.lite.TFLiteConverter转换 TensorFlow模型。TensorFlow模型是使用 SavedModel 格式存储的，并通过高阶 tf.keras.* API（Keras 模型）或低阶 tf.* API（用于生成具体函数）生成。具体来说，开发者可以使用以下三个选项转换 TensorFlow模型：

1. tf.lite.TFLiteConverter.from_saved_model()（推荐）：转换SavedModel。
2. tf.lite.TFLiteConverter.from_keras_model()：转换Keras模型。
3. tf.lite.TFLiteConverter.from_concrete_functions()：转换具体函数。

在接下来的内容中，将详细讲解上述三种转换方式的用法。

（1）转换 SavedModel（推荐）

例如在下面的代码中，例演示了将 SavedModel 转换为 TensorFlow Lite 模型的过程。

import tensorflow as tf
# 转换模型
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) # path to the SavedModel directory
tflite_model = converter.convert()
# 保存模型
with open('model.tflite', 'wb') as f:
  f.write(tflite_model)

（2）转换Keras模型

在下面的实例文件cov01.py中，演示了将Keras模型转换为TensorFlow Lite模型的过程。

import tensorflow as tf

#使用高级tf.keras.*API创建模型
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') # compile the model
model.fit(x=[-1, 0, 1], y=[-3, -1, 1], epochs=5) # train the model
# (to generate a SavedModel) tf.saved_model.save(model, "saved_model_keras_dir")

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

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

执行后会将创建的模型转换为TensorFlow Lite模型，并保存为文件model.tflite，如图13-6所示。

图13-6  TensorFlow Lite模型

（3）转换具体函数

到作者写作本书时为止，目前仅支持转换单个具体函数。例如在下面的代码中，演示了将具体函数转换为 TensorFlow Lite模型的过程。

import tensorflow as tf
#使用低级tf.*API创建模型
class Squared(tf.Module):
  @tf.function
  def __call__(self, x):
    return tf.square(x)
model = Squared()
# (ro run your model) result = Squared(5.0) # This prints "25.0"
# (to generate a SavedModel) tf.saved_model.save(model, "saved_model_tf_dir")
concrete_func = model.__call__.get_concrete_function()

#转换模型
converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
tflite_model = converter.convert()

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

注意：在开发过程中，建议大家使用上面介绍的Python API方式转换TensorFlow Lite模型。

2. 命令行工具

如果已经使用 pip安装了 TensorFlow，请按下文所示使用 tflite_convert 命令：（如果您已从源代码安装了 TensorFlow，则可以在命令行中使用如下命令转换：

tflite_convert

如果要查看所有的可用标记，请使用以下命令：

$ tflite_convert --help

`--output_file`. Type: string. Full path of the output file.
`--saved_model_dir`. Type: string. Full path to the SavedModel directory.
`--keras_model_file`. Type: string. Full path to the Keras H5 model file.
`--enable_v1_converter`. Type: bool. (default False) Enables the converter and flags used in TF 1.x instead of TF 2.x.

You are required to provide the `--output_file` flag and either the `--saved_model_dir` or `--keras_model_file` flag.

（1）转换SavedModel

将SavedModel转换为TensorFlow Lite模型的命令如下：

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

（2）转换Keras H5模型

将Keras H5模型转换为TensorFlow Lite模型的命令如下：

tflite_convert \
  --keras_model_file=/tmp/mobilenet_keras_model.h5 \
  --output_file=/tmp/mobilenet.tflite

在使用命令方式或者Python API方式转换为TensorFlow Lite模型后，接下来可以在里面添加元数据，从而在设备上部署模型时可以更轻松地创建平台专用封装容器代码。最后使用 TensorFlow Lite 解释器在客户端设备（例如移动设备、嵌入式设备）上运行模型。

13.3.2  将TensorFlow RNN转换为TensorFlow Lite

通过使用TensorFlow Lite，能够将TensorFlow RNN模型转换为TensorFlow Lite的融合 LSTM 运算。融合运算的目的是为了最大限度地提高其底层内核实现的性能，同时也提供了一个更高级别的接口来定义如量化之类的复杂转换。在TensorFlow中的RNN API的变体有很多，我们的转换方法主要包括如下两个方面：

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

（1）转换器 API

该功能是TensorFlow 2.3 版本的一部分，也可以通过 tf-nightly pip 或从头部获得。当通过 SavedModel 或直接从 Keras 模型转换到 TensorFlow Lite 时，可以使用此转换功能。例如在下面的代码中，演示了将保存的模型转换为TensorFlow Lite模型的方法。

#构建保存的模型
#此处的转换函数是对应于包含一个或多个Keras LSTM层的TensorFlow模型的导出函数
saved_model, saved_model_dir = build_saved_model_lstm(...)
saved_model.save(saved_model_dir, save_format="tf", signatures=concrete_func)
# 转换模型
converter = TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()

再看下面的代码，演示了将Keras模型转换为TensorFlow Lite模型的方法。

#建立一个Keras模型
keras_model = build_keras_lstm(...)
#转换模型
converter = TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()

在现实应用中，使用最多的是实现Keras LSTM到TensorFlow Lite的开箱即用的转换。请看下面的实例文件cov02.py，功能是使用 Keras构建用于实现MNIST识别的TFLite LSTM融合模型，然后将其转换为 TensorFlow Lite模型。

实例文件cov02.py的具体实现代码如下所示。

（1）构建MNIST LSTM模型，代码如下：

import numpy as np
import tensorflow as tf

model = tf.keras.models.Sequential([
    tf.keras.layers.Input(shape=(28, 28), name='input'),
    tf.keras.layers.LSTM(20, time_major=False, return_sequences=True),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation=tf.nn.softmax, name='output')
])
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.summary()

（2）训练和评估模型，本实例将使用MNIST数据训练模型。代码如下：

#加载MNIST数据集
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
x_train = x_train.astype(np.float32)
x_test = x_test.astype(np.float32)

# 如果要快速测试流，请将其更改为True。
# #使用小数据集和仅1个epoch进行训练。该模型将工作得很差，但这提供了一种测试转换是否端到端工作的快速方法。
_FAST_TRAINING = False
_EPOCHS = 5
if _FAST_TRAINING:
  _EPOCHS = 1
  _TRAINING_DATA_COUNT = 1000
  x_train = x_train[:_TRAINING_DATA_COUNT]
  y_train = y_train[:_TRAINING_DATA_COUNT]

model.fit(x_train, y_train, epochs=_EPOCHS)
model.evaluate(x_test, y_test, verbose=0)

（3）将Keras模型转换为TensorFlow Lite模型，代码如下：

run_model = tf.function(lambda x: model(x))
#这很重要，让我们修正输入大小。
BATCH_SIZE = 1
STEPS = 28
INPUT_SIZE = 28
concrete_func = run_model.get_concrete_function(
    tf.TensorSpec([BATCH_SIZE, STEPS, INPUT_SIZE], model.inputs[0].dtype))

#保存模型的目录
MODEL_DIR = "keras_lstm"
model.save(MODEL_DIR, save_format="tf", signatures=concrete_func)

converter = tf.lite.TFLiteConverter.from_saved_model(MODEL_DIR)
tflite_model = converter.convert()

（4）检查转换后的 TensorFlow Lite模型，现在开始加载 TensorFlow Lite 模型并使用 TensorFlow Lite Python解释器来验证结果。代码如下：

#使用TensorFlow运行模型以获得预期结果.
TEST_CASES = 10

#使用TensorFlow Lite运行模型
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

for i in range(TEST_CASES):
  expected = model.predict(x_test[i:i+1])
  interpreter.set_tensor(input_details[0]["index"], x_test[i:i+1, :, :])
  interpreter.invoke()
  result = interpreter.get_tensor(output_details[0]["index"])

  #断言TFLite模型的结果是否与TF模型一致。
  np.testing.assert_almost_equal(expected, result)
  print("Done. The result of TensorFlow matches the result of TensorFlow Lite.")

  # TfLite融合的Lstm内核是有状态的，接下来需要重置状态，即清理内部状态
  interpreter.reset_all_variables()

执行后会输出：

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
lstm (LSTM)                  (None, 28, 20)            3920      
_________________________________________________________________
flatten (Flatten)            (None, 560)               0         
_________________________________________________________________
output (Dense)               (None, 10)                5610      
=================================================================
Total params: 9,530
Trainable params: 9,530
Non-trainable params: 0
_________________________________________________________________
Epoch 1/5
1875/1875 [==============================] - 33s 17ms/step - loss: 0.3559 - accuracy: 0.8945
Epoch 2/5
1875/1875 [==============================] - 32s 17ms/step - loss: 0.1355 - accuracy: 0.9589
Epoch 3/5
1875/1875 [==============================] - 32s 17ms/step - loss: 0.0974 - accuracy: 0.9708
Epoch 4/5
1875/1875 [==============================] - 33s 17ms/step - loss: 0.0769 - accuracy: 0.9764
Epoch 5/5
1875/1875 [==============================] - 31s 17ms/step - loss: 0.0658 - accuracy: 0.9796
Done. The result of TensorFlow matches the result of TensorFlow Lite.
Done. The result of TensorFlow matches the result of TensorFlow Lite.
Done. The result of TensorFlow matches the result of TensorFlow Lite.
Done. The result of TensorFlow matches the result of TensorFlow Lite.
Done. The result of TensorFlow matches the result of TensorFlow Lite.
Done. The result of TensorFlow matches the result of TensorFlow Lite.
Done. The result of TensorFlow matches the result of TensorFlow Lite.
Done. The result of TensorFlow matches the result of TensorFlow Lite.
Done. The result of TensorFlow matches the result of TensorFlow Lite.
Done. The result of TensorFlow matches the result of TensorFlow Lite.

并且在“keras_lstm”目录中会保存创建的模型文件，如图13-7所示。

图13-7  创建的模型文件

（5）最后让检查转换后的TFLite模型，此时可以看到 LSTM将采用融合格式。如图13-8所示。

图13-8  转换后的TFLite模型

请注意，本实例创建的是融合的LSTM 操作而不是未融合的版本。本实例并不会试图将模型构建为真实世界的应用程序，而只是演示如何使用 TensorFlow Lite。大家可以使用 CNN 模型构建更好的模型。当实现Keras LSTM到 TensorFlow Lite的开箱即用转换时，强调与Keras运算定义相关的 TensorFlow Lite的LSTM协定也是十分重要的：

1. input张量的0 维是批次epoch的大小。
2. recurrent_weight张量的0维是输出的数量。
3. weight和recurrent_kernel张量进行了转置。
4. 转置后的 weight 张量、转置后的 recurrent_kernel张量，以及bias张量沿着0 维被拆分成了 4 个大小相等的张量，这些张量分别对应 input gate、forget gate、cell和output gate。