从零搭建视觉模型 15：保存模型、恢复训练与独立推理

学习初稿 · 第 15 课。 本篇已开放网页阅读，用于先建立完整认识。文中手算值和示意记录不代表真实训练成绩；真实数据训练、预训练权重和板端测试以本篇具体说明为准，本篇已附代码练习，实际验证范围见文末与包内记录。

训练进程结束以后，还剩下什么

模型存在于内存中，进程退出就需要从文件重新恢复。只打印最终准确率并没有交付可用模型；保存了参数却忘记类别顺序和输入处理，也无法保证另一个程序正确使用。

本课延续小型分类器，但先用一个更小的独立模型验证保存机制。它不依赖前一篇的文件，避免“必须先训练几个小时才能检查加载”的问题。

推理文件与恢复训练文件内容不同

推理至少需要模型结构约定、参数、类别表与预处理。恢复训练通常还需要优化器状态、调度器状态、当前轮数、混合精度缩放器状态（若使用）及必要的随机状态。

AdamW 等优化器包含历史统计，只恢复模型参数再新建优化器，虽然还能训练，却不等于精确接着原轨迹运行。若要复现批次中途的状态，还要处理采样顺序与数据迭代进度，单独保存 epoch 不够。

PyTorch 常用 state_dict 保存参数与缓冲区，并重建结构后加载；保存最佳模型时要避免后续训练继续修改所引用的张量。官方保存与加载教程

用同一个输入检查保存前后

下面代码只创建临时文件，并在退出时清理。它证明参数加载的一致性，不证明模型具备分类能力。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from torch import nn

torch.manual_seed(7)

def make_model():
    return nn.Sequential(nn.Flatten(), nn.Linear(12, 2))

model = make_model().eval()
x = torch.arange(12, dtype=torch.float32).reshape(1, 3, 2, 2) / 12
with torch.no_grad():
    expected = model(x)

with TemporaryDirectory() as directory:
    path = Path(directory) / 'checkpoint.pt'
    torch.save({'model': model.state_dict(),
                'classes': ['airplane', 'automobile'],
                'shape': [3, 2, 2]}, path)
    data = torch.load(path, map_location='cpu', weights_only=True)
    restored = make_model().eval()
    restored.load_state_dict(data['model'], strict=True)
    with torch.no_grad():
        actual = restored(x)
    torch.testing.assert_close(actual, expected)
    index = actual.argmax(1).item()
    print(data['classes'][index], actual.shape)

map_location 控制加载位置，避免把文件中记录的设备当成当前必然存在的设备。strict=True 会检查参数名称与结构匹配，先不要用忽略缺失项来掩盖版本不一致。只加载可信来源文件，并核对所用版本支持的加载选项。

类别和预处理也要保存

飞机与汽车模型应保存 RGB、输入尺寸、轴布局、数值缩放、插值策略和类别顺序。若训练加入了随机增强，推理配置通常保存确定的基础处理，而不是把随机裁剪原样带入部署。

模型本身是否包含归一化同样重要。若外部先除 255，图内又除一次，数值会缩小 255 倍。最好把“图外输入是什么，图内还做什么”写成接口表，并保存一份参考输入及输出，供不同语言实现比较。

对于前面的真实分类脚本，best.pt 保存的是推理所需参数和部分元数据，没有优化器状态，因此应称为最佳推理检查点，不能宣称支持精确恢复训练。要继续训练，需要明确补充状态并验证。

独立推理程序的边界

独立程序应能完成读取模型、读取图片、检查尺寸与通道、预处理、推理、解释类别和输出错误信息。模型加载最好只做一次，不能每处理一帧都重新读文件。

评估时调用 eval()，并使用 no_grad() 或适用的推理模式减少梯度相关开销。遇到坏文件应明确失败，不要静默用全零图片代替；这种替代会让现场问题变成看似合理但错误的分类。

输出一致性的第一步是同一个数值张量，再往外扩展到同一个图片文件。如果直接从摄像头开始比较，两次采集本来就不同，会混淆模型差异与输入差异。

练习

1. 只保存参数，换一个最后输出三类的结构，能严格加载二分类参数吗？
2. 推理结果一致，是否说明继续训练时轨迹也一致？
3. 为什么需要保存类别表而不是重新按目录排序？

下一阶段开始理解成熟网络。我们先看 ResNet 如何组织更深的计算，再学习利用预训练参数。

综合实践

想把这些知识接成一套可运行的代码，可以进入实现 MLP：从损失计算、训练到 ONNX 推理。实践包含自生成数据、训练与验证、模型保存、独立推理、ONNX 导出和数值对照；遇到问题可按实践中的索引返回对应理论章节。

配套代码练习

"""保存加载与精确恢复下一步。运行方式与练习见 README.md。"""
from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from torch import nn
torch.manual_seed(7)
torch.set_num_threads(2)
x,y = torch.randn(4,3),torch.tensor([0,1,0,1])
model = nn.Linear(3,2)
optimizer = torch.optim.SGD(model.parameters(),lr=0.1,momentum=0.9)
def step(net,opt):
    opt.zero_grad()
    loss = nn.functional.cross_entropy(net(x),y)
    loss.backward(); opt.step()
    return loss.item()
step(model,optimizer)
with TemporaryDirectory() as d:
    p = Path(d)/'checkpoint.pt'
    torch.save({'model':model.state_dict(),'optimizer':optimizer.state_dict(),
                'rng':torch.get_rng_state(),'classes':['airplane','automobile']},p)
    state = torch.load(p,weights_only=True)
    restored = nn.Linear(3,2)
    restored.load_state_dict(state['model'])
    resumed = torch.optim.SGD(restored.parameters(),lr=0.1,momentum=0.9)
    resumed.load_state_dict(state['optimizer'])
    torch.testing.assert_close(model(x),restored(x),rtol=0,atol=0)
    step(model,optimizer)
    torch.set_rng_state(state['rng'])
    step(restored,resumed)
    for a,b in zip(model.parameters(),restored.parameters()):
        torch.testing.assert_close(a,b,rtol=0,atol=0)
print('保存加载输出一致；恢复动量后下一步参数完全一致')

cd 'content/posts/AI-深度学习/从零搭建视觉模型/15-保存恢复与独立推理/code' && python3 main.py

本篇已附独立实验包，在代码卡片中查看源码、下载或复制运行命令。解压后执行 python3 main.py；依赖安装见包内 README.md。

观察目标： 固定 CPU 输入下恢复后的输出及下一次更新均一致。检查点写到临时目录；完整训练还需保存数据采样器等状态。

动手修改： 删除恢复 optimizer 的语句，观察下一步参数是否一致。为什么恢复模型权重还不够？

所有输入均在本篇代码中构造；运行范围和验证记录见包内 verification.md。

完整阅读目录 · 上一篇 · 下一篇：MLP 综合实践