第4章序列化与检查点

序列化与检查点：模型参数的保存与加载

模型训练最怕什么？我最怕训练到一半，机房断电了。或者跑了一周，发现某个超参数设错了。嗯，这时候你就知道序列化和检查点有多重要了。

说白了，序列化就是把模型的内存状态「拍个快照」存到硬盘上。检查点管理，就是怎么组织这些快照，让它们不乱、不丢、还能随时恢复。

Flax 中的参数保存与加载

Flax 模型的核心是参数。参数是一个嵌套的字典结构，里面装着每个层的权重和偏置。保存参数，本质上就是把这个字典序列化。

我个人习惯用 joblib 或者 pickle，但在生产环境中，我更推荐用 msgpack。为什么？因为 msgpack 比 pickle 快，而且跨语言兼容。

import jax.numpy as jnp
import msgpack
import flax

# 假设 model_state 是训练好的参数字典
model_state = {'params': {'Dense_0': {'kernel': jnp.ones((3, 2)), 'bias': jnp.zeros((2,))}}}

# 保存
bytes_data = msgpack.dumps(model_state)
with open('model.msgpack', 'wb') as f:
    f.write(bytes_data)

# 加载
with open('model.msgpack', 'rb') as f:
    loaded_state = msgpack.loads(f.read())

这里有个坑。我曾经在加载时发现参数形状对不上。原因是模型结构改了，但旧参数还在用。所以，保存参数时最好也保存一份模型配置。

注意：JAX 的数组是设备相关的。如果你在 GPU 上保存，在 CPU 上加载，可能会触发设备传输。建议保存前用 jax.device_get() 把数组拉到主机内存。

使用 Orbax 进行 Checkpoint 管理

Orbax 是 Google 官方推荐的检查点库。它解决了几个痛点：并发写入、异步保存、版本管理。

你想想看，分布式训练时，多个进程同时写检查点，很容易出冲突。Orbax 用文件锁和原子写入解决了这个问题。

import orbax.checkpoint as ocp

# 创建检查点管理器
checkpoint_dir = '/tmp/my_checkpoints'
options = ocp.CheckpointManagerOptions(
    max_to_keep=3,           # 最多保留3个检查点
    save_interval_steps=1000 # 每1000步保存一次
)
mgr = ocp.CheckpointManager(
    checkpoint_dir, 
    options=options
)

# 保存
step = 10000
mgr.save(step, args=ocp.args.StandardSave(model_state))

# 加载最新检查点
restored = mgr.restore(mgr.latest_step())

我个人特别喜欢 max_to_keep 这个参数。它自动删除旧的检查点，省得你手动清理。我在一个项目中训练了30万步，如果不限制，硬盘早就爆了。

小技巧：Orbax 支持异步保存。设置 async_manager=True，保存操作不会阻塞训练循环。适合大规模训练。

断点续训策略

断点续训，说白了就是训练到一半停了，从最近一次保存的地方继续跑。但这里有个细节：优化器状态也要保存。

很多新手只保存模型参数，不保存优化器状态。结果恢复训练后，学习率调度器从头开始，优化器的动量也丢了。训练效果大打折扣。

# 保存完整训练状态
train_state = {
    'params': params,
    'opt_state': opt_state,
    'step': step,
    'rng': rng
}

# 恢复时
restored = mgr.restore(mgr.latest_step())
params = restored['params']
opt_state = restored['opt_state']
step = restored['step']
rng = restored['rng']

我曾经犯过一个错误：恢复训练时忘了恢复随机数种子。结果数据打乱顺序变了，验证集评估结果对不上。嗯，从那以后，我每次都会把 rng 也存下来。

  核心原则：断点续训要保存「可恢复训练的一切」。包括参数、优化器状态、学习率调度器状态、随机数种子、当前步数。

模型版本控制

模型版本控制，不是简单的改个文件名。它要解决三个问题：

可追溯：知道每个版本用了什么数据、什么超参数
可回滚：发现新版本效果差，能快速切回旧版本
可对比：不同版本之间能比较性能差异

我建议用目录结构来管理版本：

checkpoints/
├── v1.0/
│   ├── model.msgpack
│   ├── config.yaml
│   └── metrics.json
├── v1.1/
│   ├── model.msgpack
│   ├── config.yaml
│   └── metrics.json
└── v2.0/
    ├── model.msgpack
    ├── config.yaml
    └── metrics.json

每个版本目录下，除了模型参数，还要保存配置文件和评估指标。这样你一眼就能看出哪个版本好。

Orbax 其实内置了版本管理。它用步数作为版本号，但你也可以自定义命名规则。

# 自定义版本命名
mgr.save(
    step, 
    args=ocp.args.StandardSave(train_state),
    custom_metadata={'version': 'v1.0', 'accuracy': 0.95}
)

我个人习惯在 custom_metadata 里记录验证集指标。这样不用加载模型，就能知道哪个版本最好。

生产建议：模型版本控制最好和代码版本控制联动。在 config.yaml 里记录 Git commit hash，这样模型和代码能一一对应。

知识体系总览

下面这张图总结了本章的核心逻辑。从参数序列化到检查点管理，再到断点续训和版本控制，环环相扣。

这张图展示了四个核心模块的关系。参数序列化是基础，Orbax 提供管理框架，断点续训是应用场景，版本控制是工程保障。缺一个，生产环境就容易出问题。

  一句话总结：序列化是手段，检查点是载体，断点续训是目的，版本控制是保障。四者缺一不可。

公众号:蓝海资料掘金营，微信deep3321

📚 课程目录

序列化与检查点：模型参数的保存与加载

Flax 中的参数保存与加载

使用 Orbax 进行 Checkpoint 管理

断点续训策略

模型版本控制

知识体系总览