TensorFlow 2.0 部署：模型导出

节选自《简单粗暴 TensorFlow 2.0》，回复关键字“手册”获取合集
接下来我们将介绍 TensorFlow 中模型的部署与导出，本文介绍使用 SavedModel 完整导出模型。

使用 SavedModel 完整导出模型

在部署模型时，我们的第一步往往是将训练好的整个模型完整导出为一系列标准格式的文件，然后即可在不同的平台上部署模型文件。这时，TensorFlow 为我们提供了 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])

关注公众号渣渣渣辉,回复手册获取更多.

如果帮到了你，就请收藏评论加点赞，多谢支持。

原文来自微信公众号，侵删。