基础模型checkpoint实战指南：从训练到部署的完整避坑手册

1次阅读

共计 2841 个字符，预计需要花费 8 分钟才能阅读完成。

上周同事小李的惨痛经历：在训练了 3 天的 BERT 模型即将完成时，服务器突然断电。由于没有正确设置 checkpoint 保存策略，所有中间状态全部丢失，不得不从零开始训练。更糟糕的是，当他试图加载上周保存的 checkpoint 继续微调时，发现模型性能断崖式下跌——后来发现是因为 PyTorch 版本升级导致 state_dict 结构不兼容。

这样的故事每天都在发生。checkpoint 看似简单，但暗藏玄机。下面我们就从框架差异到生产部署，彻底掌握这个模型训练中的 ” 救命稻草 ”。

功能	PyTorch 实现	TensorFlow 实现
保存完整模型	`torch.save(model.state_dict())`	`tf.saved_model.save(model)`
仅保存权重	`torch.save(model.state_dict())`	`tf.train.Checkpoint.save()`
加载权重	`model.load_state_dict()`	`tf.train.Checkpoint.restore()`
保存优化器状态	包含在 state_dict 中	需显式传入优化器对象
自定义保存内容	可 pickle 任何 python 对象	需继承 tf.train.Checkpointable

import torch
import torch.nn as nn

# 1. 定义示例模型
class DemoModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(10, 2)

    def forward(self, x):
        return self.layer(x)

model = DemoModel()
optimizer = torch.optim.Adam(model.parameters())

# 2. 模拟训练过程（带梯度累积）for epoch in range(5):
    # ... 训练代码省略...

    # 3. 保存 checkpoint（每 2 个 epoch 保存一次）if epoch % 2 == 0:
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),  # 模型参数
            'optimizer_state_dict': optimizer.state_dict(),  # 优化器状态
            'loss': 0.01,  # 自定义元数据
        }, f'checkpoint_epoch_{epoch}.pth')

# 4. 加载 checkpoint 进行热启动
checkpoint = torch.load('checkpoint_epoch_4.pth')
model.load_state_dict(checkpoint['model_state_dict'])  # 必须先初始化相同结构的模型
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
print(f"从 epoch {checkpoint['epoch']}恢复训练")

import tensorflow as tf

# 1. 创建可追踪对象
model = tf.keras.Sequential([tf.keras.layers.Dense(2, input_shape=(10,))
])
optimizer = tf.keras.optimizers.Adam()

# 2. 实例化 Checkpoint 管理器（带版本控制）checkpoint = tf.train.Checkpoint(
    model=model, 
    optimizer=optimizer
)
manager = tf.train.CheckpointManager(
    checkpoint, 
    directory='./tf_ckpts', 
    max_to_keep=3  # 只保留最近的 3 个 checkpoint
)

# 3. 训练循环
for epoch in range(5):
    # ... 训练代码省略...

    # 4. 保存 checkpoint（带 epoch 编号）if epoch % 2 == 0:
        save_path = manager.save(checkpoint_number=epoch)
        print(f"Checkpoint saved at {save_path}")

# 5. 恢复最新 checkpoint
latest_checkpoint = manager.latest_checkpoint
if latest_checkpoint:
    checkpoint.restore(latest_checkpoint)
    print(f"从 {latest_checkpoint} 恢复模型")

PyTorch DDP 模式：

# 只在主进程保存即可避免重复
if torch.distributed.get_rank() == 0:
    torch.save(checkpoint, 'shared_ckpt.pth')

TensorFlow MultiWorkerMirroredStrategy：

# 会自动处理跨 worker 同步
options = tf.train.CheckpointOptions(experimental_io_device='/job:localhost')
checkpoint.save(file_prefix, options=options)

建议结合 MLflow 进行全生命周期管理：

import mlflow

with mlflow.start_run():
    # 训练代码...
    mlflow.pytorch.log_model(
        pytorch_model=model,
        artifact_path="model",
        registered_model_name="TextClassifier"
    )
    # 自动记录超参数和指标