TensorFlow保存模型及参数

在 TensorFlow 中，可以使用 `tf.saved_model` 模块保存模型及其参数。这个模块提供了一个 API，用于将 TensorFlow 模型导出为 SavedModel 格式。SavedModel 是一种可移植的模型格式，可用于部署到不同的平台上。

以下是保存模型及其参数的步骤：

1. 定义模型并训练它。

2. 创建一个 `tf.saved_model.Builder` 对象。

import tensorflow as tf

# 定义模型
model = tf.keras.models.Sequential([
  tf.keras.layers.Dense(64, activation='relu'),
  tf.keras.layers.Dense(10)
])

# 训练模型
model.compile(optimizer=tf.keras.optimizers.Adam(0.01),
              loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
              metrics=[tf.keras.metrics.CategoricalAccuracy()])
model.fit(x_train, y_train, epochs=5)

3. 使用 `tf.saved_model.Builder` 对象保存模型及其参数。

# 创建一个 tf.saved_model.Builder 对象
export_dir = 'saved_model/1'
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)

# 定义输入和输出的 Tensor 对象
inputs = {'input': tf.saved_model.utils.build_tensor_info(model.input)}
outputs = {'output': tf.saved_model.utils.build_tensor_info(model.output)}

# 定义签名方法
signature_def = tf.saved_model.signature_def_utils.build_signature_def(
    inputs=inputs,
    outputs=outputs,
    method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)

# 将模型和签名方法添加到 SavedModelBuilder 对象中
builder.add_meta_graph_and_variables(
    sess=tf.keras.backend.get_session(),
    tags=[tf.saved_model.tag_constants.SERVING],
    signature_def_map={
        tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def
    })

# 保存模型
builder.save()

以上代码将模型及其参数保存在 `saved_model/1` 目录下。该目录中包含一个 `saved_model.pb` 文件和一个 `variables` 目录。`saved_model.pb` 文件定义了模型的计算图，`variables` 目录存储了模型的参数。

可以使用以下代码加载已保存的模型：

# 加载 SavedModel
loaded_model = tf.saved_model.load(export_dir)

# 获取模型的签名方法
signature = list(loaded_model.signatures.values())[0]

# 运行模型
result = signature(input=tf.constant([[1.0, 2.0, 3.0, 4.0]]))['output']
print(result)