MNIST手写数字识别大概是每个学深度学习的人都会做的第一个正儿八经的项目,模型结构简单、数据量小、训练速度快,拿来练手再合适不过。但很多人练到"模型能跑、准确率还行"就停下来了,模型保存和恢复这块要么顺手点一下save,要么压根没管。直到某天你训练到一半发现机房断电、程序被杀、或者调参调崩了想退回上一个版本,才发现自己根本没给模型留后路。
这篇笔记我专门把"模型的保存与恢复"这个看似简单、实际坑不少的问题拆开讲透。内容包括:模型文件里到底存了哪些东西、如何在Keras/TensorFlow里把MNIST模型完整保存下来、训练中断后怎么接续、以及恢复模型时最常见的五个翻车场景。不论你是零基础刚跑通第一个demo,还是打算把玩过的模型部署出去,这部分内容都能直接用上。
1. 模型训练中断事故与checkpoint的必要性
1.1 一个真实的"训到断电"场景
先讲个我自己经历过的场景。有一回我在本地机器上跑一个MNIST的加强版实验,加了数据增强和两层卷积,一个epoch大概要几十秒,计划跑50个epoch。跑到第23个epoch的时候,笔记本没插电源,合盖休眠了。等我再打开终端,训练进程已经被系统杀掉,控制台里只剩报错信息。
当时我心头一紧,因为代码里没有加任何自动保存逻辑,这意味着前面22个epoch的优化成果全部清零,又得从头开始跑。更气人的是,第20个epoch之后模型的验证准确率已经爬到了99%以上,就差最后几步。那一瞬间我才真正意识到,"模型保存"不是训练结束之后才考虑的事,而是从训练一开始就要规划好的基础设施。
这个场景在真实项目里太常见了,训练任务跑在云服务器上,网络抖动、资源回收、代码bug,任何一个意外都能让你的进度归零。MNIST虽然训练成本低,重来一遍无所谓,但如果你迁移到更大的数据集上,一次训练可能跑几个小时甚至几天,没有checkpoint机制的后果就是灾难性的。
1.2 模型保存到底在解决什么问题
拆开来说,模型的保存与恢复解决三个层面的问题。
第一层是断点续训。训练过程被打断后,你能从上一次保存的位置继续,而不必从头再来。这需要保存的不只是网络的权重,还包括优化器状态、当前epoch轮次、学习率调度器的进度等信息。
第二层是模型复用。训练好的模型可以被另一个脚本、另一台机器、甚至另一个团队加载使用。大家用同一个模型文件做推理、做迁移学习、做对比实验,不需要每人自己重新训练一份。
第三层是版本回溯。调参过程中经常出现"改了个学习率之后效果反而变差了"的情况,这时候如果你保存过之前效果好的版本,就能快速回退,而不是靠记忆手动改代码重新复现。
这三层需求对应着不同的保存策略,后面我会逐个展开。但不论哪种策略,本质都离不开一件事:你需要把"模型当前的状态"完整持久化到磁盘上。这个"状态"到底包含什么,就是下一节要讲的重点。
2. 模型文件里到底存了什么:不止weight这一件事
2.1 权重参数与偏置:模型"知识"的载体
从数学角度看,一个全连接层做的事情就是y = Wx + b。训练过程不断调整的就是权重矩阵W和偏置向量b。MNIST分类模型的核心能力——把784维的像素向量映射到10个类别的概率分布——完全由这些参数决定。
以我自己常用的三层全连接网络为例,输入层784个节点、隐藏层128个节点、输出层10个节点。光第一层的权重矩阵就有784乘以128个参数,加上偏置,全部参数量大概在10万级别。对这些参数的保存,本质上就是把权重矩阵和偏置向量的数值写进文件。HDF5文件、NPY文件甚至纯文本格式都能做到。
但这只是"模型的知识"这部分。
2.2 网络结构:光有权重无法推理的原因
你可能会想,既然参数都有了,推理时初始化一个同样的网络再load进去不就行了吗?问题是,你必须在加载权重之前就知道这个网络的层数、每层神经元数量、激活函数类型、初始化方式。这些信息你当然记得住,但换了一个人、换了一台机器、过了两个月,你还能保证每个人都能手动搭出一个一模一样的网络吗?
实际工程中更稳妥的做法是,保存权重的同时把网络结构也一起序列化。这样加载的时候,程序可以从文件里重建出完整的模型,而不是依赖人工去重新定义一遍结构。这也是为什么很多框架提供了"整模型保存"方案——一个文件包含结构加权重,拿来即用。
2.3 优化器状态:最容易忽视的"训练上下文"
这一条是很多零基础学习者完全不了解的。假设你在第20个epoch保存了权重,然后中断了训练。你重新加载权重、接着从第21个epoch开始训练时,会发现验证准确率不仅没继续提升,反而还有一段波动下降。
原因在于,你没有保存优化器状态。
以Adam优化器为例,它内部维护了两个额外变量:一阶动量(梯度均值)和二阶动量(梯度平方均值)。这些变量信息反映了历史梯度的走向,决定了下一步更新方向。如果你只恢复了权重,优化器所有动量都从零开始初始化,相当于一个跑了一半的运动员突然被重置了状态,当然会有一段不适应期。
在TensorFlow的Keras接口里,model.save()保存整模型时会自动包含优化器状态,但如果你使用model.save_weights()只存权重,优化器状态就会丢失。这一点在后续"续训"场景下至关重要。
2.4 三种保存粒度对比
我用一个表格整理一下不同保存方式携带的信息差异:
| 保存方式 | 网络结构 | 权重 | 优化器状态 | 适用场景 |
|---|---|---|---|---|
| 仅保存权重 | 不记录 | 包含 | 不包含 | 推理部署、迁移学习 |
| 保存架构与权重 | 包含 | 包含 | 不包含 | 加载后进行推理或重新训练 |
| 完整checkpoint | 包含 | 包含 | 包含 | 断点续训、实验回溯 |
补充一点,TensorFlow的SavedModel和Keras的.h5整模型保存都能覆盖到"完整checkpoint"这一层,但SavedModel在部署兼容性上更胜一筹,后面我会单独讲。
3. Keras/TensorFlow下MNIST模型保存与恢复的完整实操
3.1 先搭一个可复现的三层MNIST模型
完整的保存与恢复演示需要基于一个具体模型,我用最基础的三层全连接网络来写。数据集用Keras自带的MNIST,预处理就做两件事:展平成784维向量、归一化到0到1区间。
import tensorflow as tf from tensorflow import keras from tensorflow.keras import layers # 加载并预处理MNIST数据集 (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data() x_train = x_train.reshape(-1, 784).astype("float32") / 255.0 x_test = x_test.reshape(-1, 784).astype("float32") / 255.0 # 搭建三层全连接网络 def build_model(): model = keras.Sequential([ layers.Dense(128, activation="relu", input_shape=(784,)), layers.Dense(64, activation="relu"), layers.Dense(10, activation="softmax") ]) model.compile( optimizer=keras.optimizers.Adam(learning_rate=0.001), loss="sparse_categorical_crossentropy", metrics=["accuracy"] ) return model model = build_model() model.summary()这里有个细节:input_shape=(784,)必须显式指定。如果模型是在加载之后才喂第一批数据,Keras可以通过shape推断自动构建输入层。但为了保存和恢复时结构信息完整,最好在一开始就明确输入维度。
3.2 训练中自动保存checkpoint:ModelCheckpoint回调
断点续训依赖"边训练边保存"的机制。Keras里最直接的方式是使用ModelCheckpoint回调。
import os checkpoint_path = "./checkpoints/mnist_epoch_{epoch:02d}.weights.h5" checkpoint_dir = os.path.dirname(checkpoint_path) # 准备回调:每个epoch结束保存一次权重 checkpoint_callback = keras.callbacks.ModelCheckpoint( filepath=checkpoint_path, monitor="val_accuracy", save_best_only=False, save_weights_only=True, verbose=1 ) history = model.fit( x_train, y_train, batch_size=128, epochs=30, validation_data=(x_test, y_test), callbacks=[checkpoint_callback] )这段代码会为每个epoch生成一个独立的.weights.h5文件,文件名里带上了epoch编号。这样如果你想回到第18个epoch的状态,直接加载mnist_epoch_18.weights.h5就行。
我习惯把save_best_only调成True配合monitor="val_accuracy",这样只保留验证集准确率最高的那一份。如果你磁盘空间有限,这能省下不少地方;但如果训练过程中出现过拟合,还是建议间隔保存几份,防止"最佳版本"也未必是你想要的那份。
需要注意save_weights_only=True的时候,保存的只是权重,恢复时需要先手动构建同结构的模型。如果save_weights_only=False,则回调会保存整模型文件。
3.3 训练结束后的整体状态固化:整模型保存
训练结束后,通常会做一次"状态固化":把模型结构和权重打包成单一文件,供后续加载使用。
# 保存整模型,包含结构、权重、优化器状态、编译信息 model.save("./saved_models/mnist_full_model.h5") # 或者保存为TensorFlow推荐使用的SavedModel格式 model.save("./saved_models/mnist_saved_model", save_format="tf").h5格式的好处是单文件、移动方便。SavedModel格式则是一个目录,内部包含了资产文件和变量文件,更适合生产部署场景。两种方式在Keras层面加载都很简单:
loaded_model = keras.models.load_model("./saved_models/mnist_full_model.h5")load_model会根据文件里的结构信息自动重建网络,不需要你手动调用build_model(),这也是"整模型保存"最省心的地方。
3.4 断点续训的完整流程
断点续训需要做到"恢复到中断前的状态并继续"。我把完整流程拆成三步,方便你直接对照落地。
第一步,判断是否存在历史checkpoint,如果存在则加载。
import glob checkpoint_files = sorted(glob.glob("./checkpoints/mnist_epoch_*.weights.h5")) if checkpoint_files: latest_checkpoint = checkpoint_files[-1] # 从checkpoint文件名中解析已完成的epoch数 import re match = re.search(r"epoch_(\d+)", latest_checkpoint) completed_epoch = int(match.group(1)) if match else 0 # 构建模型并加载权重 model = build_model() model.load_weights(latest_checkpoint) print(f"已加载第 {completed_epoch} 个epoch的权重") else: model = build_model() completed_epoch = 0第二步,重新compile优化器并调整训练轮数。你可能会问,之前不是已经compile过了吗?但build_model()内部已经包含了compile逻辑,所以这里只需注意一点:如果使用save_weights_only=True的方式保存,优化器状态并没有被持久化,续训时动量状态会重置。要连优化器状态一起恢复,就得依赖整模型保存。
initial_epoch = completed_epoch total_epochs = 30 model.fit( x_train, y_train, batch_size=128, epochs=total_epochs, initial_epoch=initial_epoch, validation_data=(x_test, y_test), callbacks=[checkpoint_callback] )initial_epoch参数的作用是告诉Keras从哪个epoch开始继续,这样TensorBoard之类的日志插件也能正确对齐时间轴,不会从零重新计数。
第三步,验证恢复效果。续训后的模型应该表现出"接续"的特征:loss从接近中断时的值继续下降,而不是恢复到训练初期的水平。这是我用来判断是否恢复成功的核心指标。
3.5 只加载权重用于推理与迁移学习
很多场景下你并不需要续训,只想把训练好的模型拿来对单张图片做预测,或者固定住前几层做迁移学习。这时候加载权重就够了。
# 推理场景 model = build_model() model.load_weights("./checkpoints/best_model.weights.h5") import numpy as np sample = x_test[0:1] # 取一张测试图片 pred = model.predict(sample) predicted_label = np.argmax(pred, axis=1)[0] print(f"预测结果为:{predicted_label}")注意两个前提:第一个是build_model()搭出来的网络结构必须和保存权重时的结构完全一致,这是出现频率最高的报错来源;第二个是加载权重前必须先compile或至少构建出各层,不然load_weights会因为没有可对应的层级而报错。
迁移学习场景下,经常只用load_weights加载不全连接分类头之前的部分,然后冻结这些层、换一个新的输出层重新训练。这种精细操作在MNIST里体现不明显,但思路是通用的。
4. 恢复模型时我踩过的那些坑:五个高频翻车实录
4.1 重新定义模型时结构与原模型不一致
这是所有人都会至少遇到一次的报错:ValueError: Layer count mismatch,或者类似Shapes not compatible的提示。
现象:训练阶段你在模型里多加了一个Dropout层,保存了权重。恢复阶段你在build_model()里写的代码漏掉了这个Dropout层,或者把隐藏层节点数从128改成了256,然后load_weights直接崩溃。
原因很简单:权重矩阵的形状和模型中对应层的形状对不上,框架自然不会接受这种"张冠李戴"式的加载。
解决办法:如果你使用load_weights,请务必把模型构建代码原封不动地复制过来。更好一点的方案是训练结束后直接保存整模型(包含结构),恢复时用load_model,完全不依赖手工重建网络。我在实际项目中的习惯是,凡是要跨脚本、跨机器加载模型,一律用整模型或SavedModel,尽量避免人肉复制结构代码。
4.2 只加载权重续训,结果二次训练效果反而变差
这坑我在前面已经预告过了。现象是:中断前已经训练到第20个epoch,验证准确率99%左右。恢复权重后继续训练,前两三个epoch准确率反而降到了97%甚至更低,然后才慢慢爬回来。
原因就是优化器状态没有恢复。Adam这类自适应学习率优化器,内部有动量项。从头初始化意味着优化器认为梯度信息是"空白的",它会用一个比较大的学习步长去探索,这个初始探索过程会让loss和准确率出现抖动。
解决办法有两种。一种是用整模型保存方式,不只存权重,把优化器状态也一并存下来:
model.save("./checkpoints/mnist_epoch_20_full.h5") # 恢复时 model = keras.models.load_model("./checkpoints/mnist_epoch_20_full.h5") model.compile(optimizer=keras.optimizers.Adam(0.001), loss="sparse_categorical_crossentropy", metrics=["accuracy"]) model.fit(..., initial_epoch=20, epochs=30)注意,即使是从整模型加载,我仍然建议显式调用一次compile。因为不同版本的TensorFlow在保存编译参数时可能留下兼容性隐患,显式compile能确保训练配置完全可控。
4.3 工作目录变化导致加载路径失效
这个坑比较低级,但特别隐蔽。假设你在/home/user/project目录下训练,保存路径是./checkpoints/model.h5。第二天你换了个终端,当前目录变成/home/user,运行同样的加载代码,文件找不到。
虽然有经验的开发者会立刻意识到是相对路径问题,但新手经常被绕晕。我的经验是,所有涉及模型保存与加载的路径,尽量使用绝对路径,或者用os.path.join基于项目根目录构造:
import os PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__)) CHECKPOINT_DIR = os.path.join(PROJECT_ROOT, "checkpoints") MODEL_PATH = os.path.join(CHECKPOINT_DIR, "best_model.h5")这样不管你在哪个目录下发脚本,路径都不会错。
4.4 模型加载成功但预测结果全是乱码:输入预处理不一致
这个坑不在保存与加载本身,而在于你恢复模型之后喂进去的数据和训练时的数据分布不一样,报错不会出现,结果却完全不可用。
一个很常见的场景:训练时你做了归一化x / 255.0,但预测时忘了归一化,直接喂原始像素值。MNIST的像素值范围在0到255,而网络在0到1区间的数据上学到的权重模式,面对255量级的输入,激活值很容易饱和,输出变成全乱码。
另一个差异是shape。训练时输入是(batch_size, 784),预测时某张图片如果用了np.expand_dims把自己reshape成了(1, 784),一般没问题。但如果你误保留了(28, 28)或者(784, 1),框架要么直接报错,要么把每个像素当成一个样本,预测出一个长度为784的结果。
我的自查清单就三条:数据类型是不是float32、像素范围是不是0到1、shape是不是(batch, 784)。
4.5 设备一致性:GPU上保存、CPU上加载
TensorFlow在这方面其实已经做得比较宽容了。如果你的训练在GPU上完成,保存的是GPU训练得到的权重数值,加载到CPU设备上时,权重数值本身不变,一般可以正常加载和推理。
真正容易出问题的是某些操作算子在不同设备上的行为差异。比如批量归一化(BatchNormalization)层,在训练模式和推理模式下统计量的使用方式不同。如果模型里用了BN层,加载到CPU上推理时,需要确保模型处于推理模式:model.trainable = False或调用model.predict而非在tf.GradientTape里做前向计算。
MNIST这个例子里的全连接网络没有BN层,所以坑不明显。但一旦你切换到卷积加BN的经典CNN结构,这个问题就会冒出来。建议从MNIST开始就养成习惯:推理前显式标记模型状态。
5. 从checkpoint到可部署模型:格式与流程的进阶选择
5.1 HDF5整模型与SavedModel的取舍
Keras支持两种主流保存格式,我做个对比方便你选择:
| 对比项 | HDF5 (.h5) | SavedModel |
|---|---|---|
| 文件形态 | 单个文件 | 目录,包含variables/assets/saved_model.pb |
| 可读性 | 二进制,适合程序读写 | 目录结构清晰,可查看变量文件 |
| 部署兼容性 | 主要靠Keras加载 | TensorFlow Serving、TFLite均支持 |
| 版本兼容性 | 对Keras版本较敏感 | 更稳定,推荐用于生产 |
| 自定义层支持 | 加载时需额外注册自定义对象 | 同样需要处理 |
我个人的建议是:做实验用.h5,图方便省事;一旦准备部署或需要服务化,就导出SavedModel。前者的单文件在管理实验版本时很好用,后者的兼容性在跨环境时更省心。
5.2 一个完整的推理服务导出示例
假设MNIST模型已经训练完成,现在需要把它导出成可以对外提供预测服务的SavedModel格式:
model = keras.models.load_model("./saved_models/mnist_full_model.h5") # 定义导出时接收的输入签名 @tf.function def predict_pixel(input_bytes: tf.Tensor) -> tf.Tensor: # 假设传入的是经过预处理的784维向量 logits = model(input_bytes, training=False) return tf.argmax(logits, axis=-1) # 保存为SavedModel tf.saved_model.save( model, "./export/mnist_serving_model", signatures={"serving_default": predict_pixel} )实际部署时,输入往往不是向量,而是原始图片字节流。你可以在tf.function内部加入解码、缩放、归一化等预处理步骤,这样外部调用方只需要传图片数据,不需要关心模型内部细节。这一步把"数据预处理"和"模型推理"封装到了一起,是工程化意识的一个重要体现。
5.3 关于ONNX的一些联想
如果你需要把模型从一个框架迁移到另一个框架,ONNX是一个不错的选择。MNIST这个模型完全可以导出成ONNX格式,在PyTorch和TensorFlow之间转换都相对顺畅。但坦白说,对于MNIST这个规模的项目,ONNX的收益有限,反而是如果后面接触更复杂的Transformer或者检测模型,格式统一带来的便利才会真正凸显。
5.4 小实验:模型文件大小与加载时间实测
我随手跑了个测评,三层全连接网络保存出的.h5文件大约在1到2MB左右,SavedModel目录因为包含资产和变量文件,略微大一点,总体上都很小。加载时间方面,两种格式都是在零点几秒内完成,感知不到差别。
这个数据没啥惊人的,但它说明了一个道理:MNIST体量小,保存恢复的所有问题都"看起来不严重"。正因为不严重,你才应该在这个项目上把方法论练熟了——等到你开始跑BERT、跑ResNet这种动辄几百MB甚至几个GB的模型时,再学这些就晚了。
6. 给初学者的模型管理工作流建议
6.1 项目目录规划:从第一步就别乱
我建议每个深度学习项目从创建第一天就固定目录结构,哪怕只是跑MNIST练手:
project/ ├── checkpoints/ # 训练中的临时权重,按epoch命名 ├── saved_models/ # 训练完成后的固化整模型 ├── export/ # 用于部署的SavedModel或ONNX ├── logs/ # TensorBoard日志 ├── data/ # 数据集 └── src/ # 训练、推理、工具脚本目录分开的好处是,你永远不会把"待恢复的checkpoint"和"可直接部署的模型"混在一起。前者是过程文件,后者是交付物,语义不同。
6.2 保存与加载的统一封装
我习惯写一个model_io.py工具脚本,把保存、加载逻辑统一封装起来。这样训练脚本、推理脚本、评估脚本各自调用统一接口,不用重复写路径解析和格式判断。
def save_checkpoint(model, path): model.save(path) def load_checkpoint(path, custom_objects=None): return keras.models.load_model(path, custom_objects=custom_objects or {}) def save_weights(model, path): model.save_weights(path) def load_weights(model, path): model.load_weights(path)别小看这几行封装。当项目里脚本数量变多之后,统一入口意味着你只需要在一个地方修改保存策略,其余业务代码全部无感。
6.3 模型版本命名规范
我踩过一次特别蠢的坑:明明训练出了新模型,结果因为文件名重了,直接把旧模型覆盖成了新模型,而且还没留备份。从那之后我强制自己用语义化命名规则:
模型名_数据集_结构_epoch数_验证准确率_时间戳例如:
mnist_mlp_3layer_epoch30_acc99.12_20250215_1430.h5文件名里包含了足够的信息,哪怕过了一个月回头翻目录,也能一眼看出每个文件是什么来头。准确率和时间戳是我必备的两个字段,前者方便挑最优版本,后者防止覆盖。
6.4 保存与恢复的"端到端验证"习惯
最后一条建议:不要等到训练中断了才去验证恢复流程是否好用。每次保存完模型之后,马上写一个三行代码的验证脚本:从文件加载模型、对同一张图片做预测、把预测结果和内存中的原模型输出比对。
loaded_model = load_model("./saved_models/mnist_full_model.h5") assert np.allclose(loaded_model.predict(x_test[:1]), model.predict(x_test[:1]), atol=1e-5) print("加载后模型输出与原模型一致")这个断言能帮你提前发现结构不匹配、预处理遗漏、路径错误等一系列问题。养成这个习惯,比读十篇保存恢复教程都有用。
我在实际项目里逐渐体会到一个道理:模型保存与恢复,看起来像是一个"边角料"功能,但它是整个训练实验体系的保险丝。没有这个保险丝,前面所有调参、训练、等待,都可能因为一次意外而彻底崩塌。MNIST虽然小,但正好能让你用最低的成本把所有环节踩一遍。真到了后面做大模型、做多卡训练的时候,你练好的这套保存恢复方法论,就是最基础的安全网。