企业官网建设流程全解析

模型训练到一半，电脑死机，三个小时的迭代白跑；调参了一周，终于跑出当前最好的结果，却因为忘了保存，下次实验得从零开始——TensorFlow 里的“模型的保存与恢复”，本质上就是给训练过程买个保险。但很多刚开始接触的人会在这一块踩坑，像是保存完加载报错，或者模型文件明明生成了，推理时却提示找不到变量。这篇文章主要针对 TensorFlow 2.x 的常用保存恢复方式，结合我在实际项目里验证过的做法，把原理、代码和避坑点一次性说清楚，适合正在训练模型、准备部署或者需要断点续训的人参考。

1. 核心思路：先搞清楚保存的到底是什么

1.1 从“变量”和“图结构”说起

TensorFlow 里的模型，本质上由两部分组成：一部分是“计算图”，描述了从输入到输出的数据流向和计算逻辑；另一部分是“变量”，也就是网络里的权重、偏置这些参数值。训练的过程，就是不断调整这些变量的过程。

所以保存模型，就要么“把图结构和变量一起打包”，要么“只把变量数值记下来，图结构留到你写代码时再重新定义”。这两种思路分别对应了 TensorFlow 里的 SavedModel 和 Checkpoint。很多人搞混的就是这里：Checkpoint 保存的是变量字典，它本身没有模型结构信息；而 SavedModel 是完整的服务化产物，部署推理时基本都用它。

1.2 对应到实际场景的需求拆解

你在不同阶段，用到的保存恢复方式是不同的：

• 训练中途因意外中断，需要从最近一次保存点继续训练。这时用 Checkpoint 或者 Keras 的 ModelCheckpoint 回调最合适。
• 模型训练完毕，需要放到服务器上跑预测，或者用 TensorFlow Serving 部署。这时导出成 SavedModel。
• 只想要一组权重数值，用于迁移学习或手动传给其他框架。这时可以用 HDF5（.h5）或 SavedModel 里的变量部分。

理解这个需求分层之后，再去看 API 就不会混乱了。

2. 技术路径：两种主流保存方案的选型

2.1 Checkpoint：训练过程的“断点续训”

Checkpoint 的核心对象是tf.train.Checkpoint，它管理的是一组变量。你可以把它理解成一个“变量管家”，告诉它“你帮我看着这些变量”，然后在需要的时候让它一次性写进磁盘。

在实际项目里，比较推荐直接配合 Keras 的ModelCheckpoint回调使用，尤其是做图像分类、文本分类这类常规训练任务时，代码量最少：

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy']) checkpoint_callback = tf.keras.callbacks.ModelCheckpoint( filepath='model_checkpoint.weights.h5', save_weights_only=True, save_best_only=True, monitor='val_loss', mode='min', verbose=1 ) model.fit(train_dataset, validation_data=val_dataset, epochs=30, callbacks=[checkpoint_callback])

关键参数解读一下：

• filepath：保存路径，可以包含模板变量比如model_{epoch:02d}.weights.h5，这样每个 epoch 都会生成独立文件。
• save_weights_only=True：只保存权重，不保存整个模型结构。恢复时你需要先构建一个一模一样的模型，再load_weights进去。
• save_best_only=True+monitor='val_loss'+mode='min'：意思是只在验证损失比历史最优更小时覆盖文件，避免磁盘被一堆中间文件塞满。

如果你需要同时保存优化器状态，实现“从上次的学习率、动量状态完全恢复”，可以把save_weights_only改为False，或者直接用tf.train.Checkpoint把所有对象传进去：

checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer) checkpoint_manager = tf.train.CheckpointManager( checkpoint, directory='./tf_ckpts', max_to_keep=3 ) # 训练循环中定期执行 checkpoint_manager.save()

CheckpointManager的max_to_keep参数很有用，它帮我保留了最近 3 次快照，同时自动清理旧的 checkpoint 文件。这样即使某次保存的文件本身是坏的，也能回退到上一个可用版本。

2.2 SavedModel：部署和分发的“标准件”

SavedModel 是 TensorFlow 官方推荐的跨平台格式，它把计算图、变量、签名信息打包成一个目录。目录里通常有assets/、variables/、saved_model.pb这些内容。训练完的模型如果打算用 TF Serving 上线，或者转成 TensorFlow Lite 跑移动端，导成 SavedModel 是稳妥路线。

导出方式非常直接：

model.save('exported_model', save_format='tf')

恢复使用也简单：

loaded_model = tf.keras.models.load_model('exported_model')

这里要说一下，model.save在 TensorFlow 2.x 中，默认格式就是 SavedModel，除非你明确指定save_format='h5'。所以上面的代码里我写了save_format='tf'，纯粹是为了显式表达意图，防止团队里其他人误以为它保存的是 H5。

SavedModel 文件在磁盘上是一个目录，不是单个文件。如果你要拷贝到其他机器，记得把整个目录一起拷走，别只拿一个saved_model.pb，否则加载时大概率报“missing variables”之类的错误。

2.3 直接保存整个模型（H5 格式）

早在 Keras 时代就流行把整个模型存成.h5，这个格式同时保存了结构、权重和训练配置。对小规模项目或快速验证来说很方便，加载时不需要重新定义模型结构：

model.save('my_model.h5') restored_model = tf.keras.models.load_model('my_model.h5')

但有个缺点得知道：H5 在保存自定义层、自定义损失函数时比较容易出兼容问题。比如你写了一个带@tf.function装饰的逻辑，或者实现了自定义tf.keras.layers.Layer子类，加载时通常会报“unable to deserialize”的错误，需要传入custom_objects参数，或者在保存时使用 SavedModel 格式。

3. 实操：从零到一完整演示保存与恢复

3.1 搭建一个可复现的示例环境

为了不引入额外变量，就用 TensorFlow 自带的 Fashion MNIST 数据集来演示。这个数据集是 28x28 的灰度图，共 10 个类别，训练集 60000 张，测试集 10000 张。网络结构也不用太复杂，三层全连接就能达到可接受的效果。

先定义模型：

import tensorflow as tf from tensorflow.keras import layers def create_model(): model = tf.keras.Sequential([ layers.Input(shape=(28, 28)), layers.Flatten(), layers.Dense(128, activation='relu'), layers.Dense(64, activation='relu'), layers.Dense(10, activation='softmax') ]) return model

之所以把模型定义写成一个函数，是为了恢复权重时能稳定地重建“一模一样的模型结构”。这一步看着不起眼，但无数报错都源于“训练时的模型和恢复时的模型结构不一致”。

数据加载和预处理：

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data() x_train = x_train.astype('float32') / 255.0 x_test = x_test.astype('float32') / 255.0

标准化是必须的，不归一化的话，全连接网络收敛速度会很惨，甚至可能出现梯度爆炸。

3.2 训练过程中使用 Checkpoint 实现断点续训

下面这段代码是训练环节的重点。我会同时使用ModelCheckpoint和EarlyStopping，前者负责保存，后者负责在验证集不再提升时提前结束训练，节省时间：

model = create_model() model.compile( optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), loss='sparse_categorical_crossentropy', metrics=['accuracy'] ) ckpt_callback = tf.keras.callbacks.ModelCheckpoint( filepath='best_model.weights.h5', save_weights_only=True, save_best_only=True, monitor='val_accuracy', mode='max' ) early_stop = tf.keras.callbacks.EarlyStopping( monitor='val_loss', patience=5, restore_best_weights=True ) history = model.fit( x_train, y_train, batch_size=128, epochs=20, validation_split=0.2, callbacks=[ckpt_callback, early_stop], verbose=1 )

validation_split=0.2的意思是自动从训练集尾部取出 20% 当验证集。注意它是在洗牌之前划分的，如果你的数据本身有顺序分布偏向，最好手动做train_test_split而不是依赖这个参数。

我在真实项目里经常还会在回调列表里加一个ReduceLROnPlateau，在验证损失平台期自动把学习率降为原来的 0.5 倍，配合EarlyStopping能多榨出一些精度。但这跟保存恢复主题关系不大，这里就不展开。

3.3 模拟“意外中断”并从检查点恢复

现在模拟一个场景：训练还没跑完，进程就被意外终止了。如果我们没有保存，一切都要重来。但现在best_model.weights.h5已经记录了历史最优权重。恢复的方法是先重建模型，再把权重加载进去：

model = create_model() model.compile( optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), loss='sparse_categorical_crossentropy', metrics=['accuracy'] ) model.load_weights('best_model.weights.h5')

这里有个经常被忽视的细节：如果你在使用ModelCheckpoint时设置了save_weights_only=True，那么恢复时就只需要load_weights。但如果保存的是完整模型（.h5或 SavedModel），恢复时就要用load_model，并且不需要提前create_model()。这两种方式别混淆，否则你会反复看见各种 shape mismatch 或者 “model not built” 的报错。

恢复之后，可以继续训练：

model.fit( x_train, y_train, batch_size=128, epochs=5, validation_split=0.2, callbacks=[ckpt_callback, early_stop], verbose=1 )

这样做的意义在于，你不需要从头开始跑完全部 epoch，而是从历史最优状态继续。如果之前的训练已经过拟合，继续训练可能反而更差，所以在恢复后先评估一下当前模型在验证集上的表现，再决定要不要继续。

3.4 推理阶段恢复模型并评估

权重恢复完成后，直接上测试集评估：

test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0) print(f"Test accuracy: {test_acc:.4f}")

假设之前训练到验证准确率大概 0.89，恢复出来的模型在测试集上也应该得到接近的结果。如果差太多，说明权重文件与模型结构对不上，或者保存和恢复时的预处理流程不一致。

用恢复好的模型做单张图片预测：

import numpy as np sample = x_test[0] sample_batch = np.expand_dims(sample, axis=0) predictions = model.predict(sample_batch) predicted_class = np.argmax(predictions, axis=1)[0] print(f"Predicted class: {predicted_class}")

单张图片必须扩充成 batch 维度，因为模型默认接收(None, 28, 28)的输入，直接传(28, 28)会报维度错误。

3.5 导出 SavedModel 用于部署

在模型表现满意之后，导出成 SavedModel：

model.save('fashion_mnist_savedmodel')

导出的目录结构大致如下：

fashion_mnist_savedmodel/ ├── assets/ ├── variables/ │ ├── variables.data-00000-of-00001 │ └── variables.index └── saved_model.pb

加载测试：

loaded_model = tf.keras.models.load_model('fashion_mnist_savedmodel') test_loss, test_acc = loaded_model.evaluate(x_test, y_test, verbose=0) print(f"Loaded model test accuracy: {test_acc:.4f}")

注意load_model加载的是整个模型，所以不需要重新构建网络结构。它甚至能保留优化器状态，如果之前编译过的话。

3.6 关于@tf.function对保存恢复的影响

TensorFlow 2.x 的默认模式是 Eager Execution（动态图），也就是每行代码立即执行。但为了性能，我们经常用@tf.function把 Python 函数转成静态图。这个机制会引入“具体函数缓存”，对保存恢复有一定影响。

比如你自定义了一个推理函数：

class MyModel(tf.keras.Model): def __init__(self): super().__init__() self.dense = layers.Dense(10, activation='softmax') @tf.function def call(self, inputs): return self.dense(inputs)

在恢复时，如果输入张量的 shape 改变了（比如从训练时的(None, 28, 28)变成推理时的(None, 784)），会触发重新 tracing。大多数情况下这是透明的，但如果你在函数里写了一些依赖具体 shape 的逻辑，比如tf.reshape(inputs, [-1, 28, 28])，shape 变了就会报错或者行为不符合预期。我的建议是：自定义层或模型里的 shape 处理尽量写在build或__init__里，不要在call里硬编码。

4. 排查手记与常见问题

4.1 加载权重时报错“shape mismatch”

这是最高频的报错之一。原因几乎总是训练时的模型结构和加载权重时的模型结构不一致。比如训练时Dense(128)后面接Dense(10)，恢复时改成了Dense(64)再接Dense(10)，权重维度对不上。

排查方法：

1. 打印模型结构：model.summary()，肉眼检查每一层的输出维度。
2. 确认输入 shape 一致：Input(shape=(28, 28))不能恢复时变成Input(shape=(784,))，否则权重虽然能加载，但第一个 Dense 层的输入矩阵计算会出问题。
3. 检查是否真的有变量名冲突：如果用了混合精度策略tf.keras.mixed_precision.set_global_policy('mixed_float16')，保存和加载时的策略不同，也可能出现变量类型不一致。尽量在加载前就设置好同样的策略。

4.2 加载 SavedModel 时报错“SavedModel not found”

这个通常是路径问题。SavedModel 是一个目录，load_model要指向那个目录本身，而不是目录里的saved_model.pb。我见过有人写成：

model = tf.keras.models.load_model('exported_model/saved_model.pb')

这是错的。应该写：

model = tf.keras.models.load_model('exported_model')

另外，如果目录是通过压缩包传输的，解压时可能丢失了隐藏文件或权限设置。建议用tar -czf打包整个目录，确保符号链接和文件权限都保留。

4.3 恢复后继续训练，loss 反而异常升高

4.3 恢复后继续训练，loss 反而异常升高

这个现象我也踩过。原本以为是权重保存坏了，后来发现是优化器状态丢失导致的。只保存权重（save_weights_only=True）时，Adam 的动量项、方差项全部没有保存。恢复训练时，优化器相当于从零开始，一开始的更新步长会偏大，loss 就可能异常波动。

解决办法：

• 如果确定要断点续训，使用tf.train.Checkpoint把优化器也一起保存，或者ModelCheckpoint设置save_weights_only=False，这会生成一个完整模型文件，内部包含优化器变量。
• 恢复后先用较小的学习率跑一个 epoch，观察 loss 是否正常下降。
• 不要拿恢复后的模型与之前记录的 loss 曲线做直接对比，因为优化器内部状态不同，训练轨迹必然不同。

4.4 自定义层导致保存恢复失败

自定义层/模型在保存时，需要保证类定义是可被 Python 找到的。也就是说，加载模型的脚本里必须 import 或重新定义同一个类。如果类名对不上，Keras 会报 “Unknown layer”。

一个常见的绕行方案是，在加载时传入custom_objects：

model = tf.keras.models.load_model( 'custom_model.h5', custom_objects={'MyLayer': MyLayer} )

但这个方法治标不治本。最佳实践其实是从一开始就把自定义层放在一个独立模块里，比如custom_layers.py，训练脚本和推理脚本都从同一个模块导入。这样既避免了类名拼写不一致，也方便别人复用。

4.5 常见问题速查表

5. 经验总结：把这些实践写进你的项目规范

5.1 设计阶段就规划好保存策略

我自己的习惯是：在项目一开始就建立一个checkpoints/目录和一个exported/目录，前者存放训练中间产物，后者存放最终部署产物。同时会在训练脚本里加一个命令行参数--restore_from，允许从某个已有 checkpoint 恢复训练。

python train.py --restore_from checkpoints/model-10.ckpt

这样不同实验之间的对比、回滚都变得非常简单。你不必每次手工修改代码，只要在启动命令时指定路径。

5.2 保存时附加元信息

检查点编号、epoch、验证指标这些信息看起来零碎，但关键时刻能救命。我通常会把它们写进模型保存路径里：

checkpoint_callback = tf.keras.callbacks.ModelCheckpoint( filepath='checkpoints/model_{epoch:02d}_valacc_{val_accuracy:.4f}.weights.h5', save_weights_only=True, save_best_only=False, monitor='val_accuracy', mode='max' )

这样文件列表本身就形成了一条“训练记录”，不用额外打开 TensorBoard 也能快速判断哪个文件是最优版本。需要注意的是，如果路径中的指标值变化过于频繁，磁盘会快速累积文件，所以这种命名方式最好搭配定期清理脚本使用。

5.3 用 SavedModel 做“最终交付物”

训练阶段的中间产物可以用 Checkpoint 或 H5，但进入交付环节，一律导出 SavedModel。原因有三：第一，SavedModel 自带签名信息，能被 TF Serving 直接识别；第二，它和 TensorFlow Lite 转换工具兼容性更好；第三，SavedModel 目录结构清晰，variables和assets分开放置，便于版本管理。

我在实际部署时，通常会在导出前先冻结计算图并做一次完整的推理验证，确保模型在“导出后”和“训练时”的输出一致。方法是取 100 个测试样本，比较model.predict()和loaded_model.predict()的结果，差异超出容忍范围就排查预处理流程。

5.4 定期验证“恢复流程”本身

有一个很容易被忽略的风险：保存代码从来不会被测试，直到某天真的需要恢复。结果是恢复时发现各种问题，而你根本没有退路。所以我现在有个习惯，每次训练脚本大改后，会主动中断一次训练，然后跑恢复脚本，看能不能从最近一个 checkpoint 接着训练。这个流程只花三分钟，但能换来很大的确定性。

尤其是当你换了机器、换了 TensorFlow 版本、或者用了不同的硬件（例如从 GPU 切到 CPU），跨环境恢复往往会出现细微问题。提前验证过，心里才有底。恢复时遇到变量不匹配，先考虑是不是 TF 版本差异导致权重格式变化，不要急着改模型结构。

5.5 关于跨框架复用的一点个人看法

如果你只是想拿走网络里的卷积核权重，用它做可视化或者喂给其他框架，用model.get_weights()拿到 numpy 数组直接操作是更轻量的方式。这种方式绕开了 TensorFlow 的序列化逻辑，是最底层的权重复刻。反过来说，如果别人只给了你一个权重文件，没有给网络结构，那你几乎是无法恢复模型的。所以“保存模型”的最小完整单位，永远要包含网络结构，或者至少包含能重建结构的代码。

我在实际使用中体会到，TensorFlow 的保存恢复机制功能本身很完整，但设计上还是偏向“完整生态”使用者。如果你只是在做小而美的实验，完全可以用更简单的 H5；如果你要走上线，那就老老实实导出 SavedModel。两种路线各有适用范围，关键是别在错误的场景里用错误的格式。希望上面这些实操路线和踩坑记录能帮你省下一些排查的时间，也欢迎你按自己项目的具体需求去调整保存策略，找到最适合自己的那一套。