如何保存TensorFlow模型？基于官方CNN教程的模型保存问询

嘿，我看你跟着官方的CNN教程搭建模型，现在想要保存它对吧？这里给你两种适配你代码风格的实用方法，都是TensorFlow里常用的保存方式：

这种格式是TensorFlow的标准导出格式，既可以用来加载继续训练，也能直接部署到TensorFlow Serving或者其他平台。

首先，我们需要先完善你的模型函数，在预测模式下添加导出签名（这样TensorFlow能识别模型的输入输出），然后训练完成后调用导出方法：

完善后的模型函数 #

import numpy as np
import tensorflow as tf
tf.logging.set_verbosity(tf.logging.INFO)

def cnn_model_fn(features, labels, mode):
    # 你的基础CNN层（补全你没写完的部分）
    input_layer = tf.reshape(features["x"], [-1, 28, 28, 1])
    conv1 = tf.layers.conv2d(
        inputs=input_layer,
        filters=32,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)
    
    # 后续的卷积、池化、全连接层（按教程补全）
    conv2 = tf.layers.conv2d(
        inputs=pool1,
        filters=64,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)
    pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)
    
    pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
    dense = tf.layers.dense(inputs=pool2_flat, units=1024, activation=tf.nn.relu)
    dropout = tf.layers.dropout(
        inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
    
    logits = tf.layers.dense(inputs=dropout, units=10)
    
    # 定义预测输出
    predictions = {
        "classes": tf.argmax(input=logits, axis=1),
        "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
    }
    
    # 关键：添加导出签名，让SavedModel知道如何接收输入
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(
            mode=mode, 
            predictions=predictions,
            export_outputs={
                'predict': tf.estimator.export.PredictOutput(predictions)
            })
    
    # 损失、训练、评估逻辑（和教程一致）
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
    
    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
        train_op = optimizer.minimize(
            loss=loss,
            global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
    
    eval_metric_ops = {
        "accuracy": tf.metrics.accuracy(
            labels=labels, predictions=predictions["classes"])}
    return tf.estimator.EstimatorSpec(
        mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)

训练后导出模型 #

当你完成模型训练后，执行下面的代码就能把模型保存成SavedModel：

# 创建Estimator实例（指定model_dir用来存检查点）
mnist_classifier = tf.estimator.Estimator(
    model_fn=cnn_model_fn, model_dir="/tmp/mnist_convnet_checkpoints")

# 这里是你的训练代码（按教程写的训练逻辑，比如train_input_fn等）
# ...（省略训练步骤）

# 导出SavedModel到指定目录
export_dir = "/tmp/saved_mnist_cnn_model"
saved_model_path = mnist_classifier.export_saved_model(
    export_dir,
    # 定义服务时的输入格式，和你的输入特征匹配
    serving_input_receiver_fn=lambda: tf.estimator.export.build_parsing_serving_input_receiver_fn(
        tf.feature_column.make_parse_example_spec([tf.feature_column.numeric_column("x", shape=[28,28])]))()
)

print(f"模型已成功保存到: {saved_model_path}")

如果你只是想保存训练过程中的权重，方便后续继续训练，那tf.estimator其实会自动帮你做这件事——只要你在创建Estimator时指定了model_dir，训练过程中它会定期把检查点（.ckpt文件）保存到这个目录里。

下次要继续训练时，只需要用同一个model_dir创建Estimator实例，它会自动加载之前保存的权重，接着上次的进度训练。

内容的提问来源于stack exchange，提问作者Abin John Thomas