【Tensorflow教程笔记】TensorFlow 模型导出

Tensorflow教程笔记

1. 基础
   TensorFlow 基础
   TensorFlow 模型建立与训练
   基础示例：多层感知机（MLP）
   卷积神经网络（CNN）
   循环神经网络（RNN）
   深度强化学习（DRL）
   Keras Pipeline
   自定义层、损失函数和评估指标
   常用模块 tf.train.Checkpoint ：变量的保存与恢复
   常用模块 TensorBoard：训练过程可视化
   常用模块 tf.data ：数据集的构建与预处理
   常用模块 TFRecord ：TensorFlow 数据集存储格式
   常用模块 tf.function ：图执行模式
   常用模块 tf.TensorArray ：TensorFlow 动态数组
   常用模块 tf.config：GPU 的使用与分配

2. 部署
   TensorFlow 模型导出
   TensorFlow Serving
   TensorFlow Lite

3. 大规模训练与加速
   TensorFlow 分布式训练
   使用 TPU 训练 TensorFlow 模型

4. 扩展
   TensorFlow Hub 模型复用
   TensorFlow Datasets 数据集载入

5. 附录
   强化学习基础简介

---

为了将训练好的机器学习模型部署到各个目标平台（如服务器、移动端、嵌入式设备和浏览器等），我们的第一步往往是将训练好的整个模型完整导出（序列化）为一系列标准格式的文件。在此基础上，我们才可以在不同的平台上使用相对应的部署工具来部署模型文件。TensorFlow 提供了统一模型导出格式 SavedModel，使得我们训练好的模型可以以这一格式为中介，在多种不同平台上部署，这是我们在 TensorFlow 2 中主要使用的导出格式。同时，基于历史原因，Keras 的 Sequential 和 Functional 模式也有自有的模型导出格式，我们也一并介绍。

使用 SavedModel 完整导出模型

在前节中我们介绍了 Checkpoint，它可以帮助我们保存和恢复模型中参数的权值。而作为模型导出格式的 SavedModel 则更进一步，其包含了一个 TensorFlow 程序的完整信息：不仅包含参数的权值，还包含计算的流程（即计算图）。当模型导出为 SavedModel 文件时，无须模型的源代码即可再次运行模型，这使得 SavedModel 尤其适用于模型的分享和部署。后文的 TensorFlow Serving（服务器端部署模型）、TensorFlow Lite（移动端部署模型）以及 TensorFlow.js 都会用到这一格式。

Keras 模型均可方便地导出为 SavedModel 格式。不过需要注意的是，因为 SavedModel 基于计算图，所以对于使用继承 tf.keras.Model 类建立的 Keras 模型，其需要导出到 SavedModel 格式的方法（比如 call ）都需要使用 @tf.function 修饰（ @tf.function 的使用方式见 前文 ）。然后，假设我们有一个名为 model 的 Keras 模型，使用下面的代码即可将模型导出为 SavedModel：

tf.saved_model.save(model, "保存的目标文件夹名称")

在需要载入 SavedModel 文件时，使用

model = tf.saved_model.load("保存的目标文件夹名称")

即可。

对于使用继承 tf.keras.Model 类建立的 Keras 模型 model ，使用 SavedModel 载入后将无法使用 model() 直接进行推断，而需要使用 model.call() 。

以下是一个简单的示例，将 前文 MNIST 手写体识别的模型 进行导出和导入。

导出模型到 saved/1 文件夹：

import tensorflow as tf
from zh.model.utils import MNISTLoader

num_epochs = 1
batch_size = 50
learning_rate = 0.001

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(100, activation=tf.nn.relu),
    tf.keras.layers.Dense(10),
    tf.keras.layers.Softmax()
])

data_loader = MNISTLoader()
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
    loss=tf.keras.losses.sparse_categorical_crossentropy,
    metrics=[tf.keras.metrics.sparse_categorical_accuracy]
)
model.fit(data_loader.train_data, data_loader.train_label, epochs=num_epochs, batch_size=batch_size)
tf.saved_model.save(model, "saved/1")

将 saved/1 中的模型导入并测试性能：

import tensorflow as tf
from zh.model.utils import MNISTLoader

batch_size = 50

model = tf.saved_model.load("saved/1")
data_loader = MNISTLoader()
sparse_categorical_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
num_batches = int(data_loader.num_test_data // batch_size)
for batch_index in range(num_batches):
    start_index, end_index = batch_index * batch_size, (batch_index + 1) * batch_size
    y_pred = model(data_loader.test_data[start_index: end_index])
    sparse_categorical_accuracy.update_state(y_true=data_loader.test_label[start_index: end_index], y_pred=y_pred)
print("test accuracy: %f" % sparse_categorical_accuracy.result())

输出:

test accuracy: 0.952000

使用继承 tf.keras.Model 类建立的 Keras 模型同样可以以相同方法导出，唯须注意 call 方法需要以 @tf.function 修饰，以转化为 SavedModel 支持的计算图，代码如下：

class MLP(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.flatten = tf.keras.layers.Flatten()
        self.dense1 = tf.keras.layers.Dense(units=100, activation=tf.nn.relu)
        self.dense2 = tf.keras.layers.Dense(units=10)

    @tf.function
    def call(self, inputs):         # [batch_size, 28, 28, 1]
        x = self.flatten(inputs)    # [batch_size, 784]
        x = self.dense1(x)          # [batch_size, 100]
        x = self.dense2(x)          # [batch_size, 10]
        output = tf.nn.softmax(x)
        return output

model = MLP()
...

模型导入并测试性能的过程也相同，唯须注意模型推断时需要显式调用 call 方法，即使用：

y_pred = model.call(data_loader.test_data[start_index: end_index])

Keras 自有的模型导出格式

由于历史原因，我们在有些场景也会用到 Keras 的 Sequential 和 Functional 模式的自有模型导出格式（H5）。我们以 keras 模型训练和保存为例进行讲解，如下是 keras 官方的 mnist 模型训练样例。

源码地址:

https://github.com/keras-team/keras/blob/master/examples/mnist_cnn.py

以上代码，是基于 keras 的 Sequential 构建了多层的卷积神经网络，并进行训练。

为了方便起见可使用如下命令拷贝到本地:

curl -LO https://raw.githubusercontent.com/keras-team/keras/master/examples/mnist_cnn.py

然后，在最后加上如下一行代码（主要是对 keras 训练完毕的模型进行保存）:

model.save('mnist_cnn.h5')

在终端中执行 mnist_cnn.py 文件，如下:

python mnist_cnn.py

该过程需要连接网络获取 mnist.npz 文件（https://s3.amazonaws.com/img-datasets/mnist.npz），会被保存到 $HOME/.keras/datasets/ 。如果网络连接存在问题，可以通过其他方式获取 mnist.npz 后，直接保存到该目录即可。

执行过程会比较久，执行结束后，会在当前目录产生 mnist_cnn.h5 文件（HDF5 格式），就是 keras 训练后的模型，其中已经包含了训练后的模型结构和权重等信息。

在服务器端，可以直接通过 keras.models.load_model("mnist_cnn.h5") 加载，然后进行推理；在移动设备需要将 HDF5 模型文件转换为 TensorFlow Lite 的格式，然后通过相应平台的 Interpreter 加载，然后进行推理。