3. PyTorch Checkpoint API 详解：torch.save() 与 torch.load()、state_dict 机制、模型与优化器的序列化

好，咱们进入实操环节。前面聊了 checkpoint 是什么、为什么重要，现在直接上手 PyTorch 的 API。说实话，torch.save() 和 torch.load() 看起来简单，但我在生产环境里见过太多人用错——模型没恢复、优化器状态丢了、跨设备加载报错……这些坑我基本都踩过一遍。

3.1 state_dict：一切序列化的核心

先搞清楚一个概念：state_dict。说白了，它就是 PyTorch 里所有可学习参数的「快照字典」。每个 nn.Module 和 optim.Optimizer 都有这个属性。

核心理解：state_dict 是一个 Python 字典，key 是参数名称（比如 conv1.weight），value 是对应的 Tensor。它不保存网络结构，只保存参数数值。

我刚开始用 PyTorch 时有个误解，以为保存模型就是保存整个网络对象。后来发现，保存 state_dict 才是标准做法。为什么？因为只保存参数，加载时更灵活——你可以换网络结构、换设备，只要参数名匹配就行。

import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的网络
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 2)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = SimpleNet()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 查看 state_dict 的结构
print("模型 state_dict 的 keys:")
for key in model.state_dict().keys():
    print(f"  - {key}")

print("\n优化器 state_dict 的 keys:")
for key in optimizer.state_dict().keys():
    print(f"  - {key}")

输出会告诉你，模型 state_dict 里是 fc1.weight、fc1.bias 这些参数名；优化器 state_dict 里则包含 state（动量缓存等）、param_groups（超参数配置）。

3.2 torch.save()：不只是保存模型

torch.save() 底层用的是 Python 的 pickle。这意味着它能序列化几乎所有 Python 对象——模型、优化器、字典、列表，甚至自定义类实例。

我个人习惯把 checkpoint 组织成一个字典，而不是只保存模型。这样一次保存，恢复时什么都有。

# 训练过程中的某个 checkpoint 保存点
checkpoint = {
    'epoch': 10,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': 0.523,
    'best_accuracy': 0.892,
    'config': {
        'batch_size': 32,
        'learning_rate': 0.001,
        'model_arch': 'SimpleNet'
    }
}

torch.save(checkpoint, 'checkpoint_epoch_10.pt')
print("Checkpoint 已保存")

文件命名建议：我习惯用 .pt 或 .pth 后缀。文件名里带上 epoch 或 loss 信息，方便后续筛选。比如 model_epoch_50_acc_0.95.pt。

3.3 torch.load()：恢复的学问

加载 checkpoint 时，有个关键参数 map_location。这个参数决定了张量被加载到哪个设备上。我在项目中遇到过最典型的错误就是：模型在 GPU 上训练，加载时没指定 map_location，结果在 CPU 机器上直接报错。

# 加载 checkpoint
checkpoint = torch.load('checkpoint_epoch_10.pt', map_location='cpu')

# 恢复模型
model = SimpleNet()
model.load_state_dict(checkpoint['model_state_dict'])

# 恢复优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

# 恢复训练状态
start_epoch = checkpoint['epoch'] + 1
best_acc = checkpoint['best_accuracy']

print(f"从 epoch {checkpoint['epoch']} 恢复训练，最佳准确率: {best_acc:.3f}")

注意：加载优化器 state_dict 前，必须先创建优化器实例。因为优化器的 load_state_dict() 会校验参数组结构是否匹配。我曾经因为先加载后创建，导致参数组不匹配而报错。

3.4 模型与优化器的序列化对比

模型和优化器的序列化，虽然都用 state_dict，但内容完全不同。我整理了一张对比表，方便你理解：

对比维度 模型 state_dict 优化器 state_dict
内容 所有可学习参数（权重、偏置） 超参数 + 动量/自适应学习率缓存
大小 较大（取决于参数量） 较小（通常几 KB 到几 MB）
恢复要求 需要相同网络结构 需要相同优化器类型和参数组
跨设备 支持（通过 map_location） 支持（但需注意设备一致性）
典型用途 推理、迁移学习、恢复训练 恢复训练（保持学习率调度状态）

嗯，这里要特别强调一点：优化器 state_dict 里的缓存是训练进度的关键。比如 Adam 优化器会保存每个参数的动量一阶矩和二阶矩估计。如果你只恢复模型而不恢复优化器，训练会从「冷启动」开始，学习率调度也会重置。这在长训练任务中可能导致收敛变慢。

3.5 完整 checkpoint 保存与恢复流程

下面是一个我在实际项目中使用的完整模板。它涵盖了保存、恢复、设备处理等常见场景。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

def save_checkpoint(model, optimizer, epoch, loss, accuracy, filepath):
    """保存完整 checkpoint"""
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss,
        'accuracy': accuracy
    }
    torch.save(checkpoint, filepath)
    print(f"Checkpoint saved to {filepath}")

def load_checkpoint(filepath, model, optimizer=None, device='cpu'):
    """加载 checkpoint，支持选择性恢复优化器"""
    checkpoint = torch.load(filepath, map_location=device)
    
    model.load_state_dict(checkpoint['model_state_dict'])
    
    if optimizer is not None:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    
    return checkpoint

# 使用示例
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleNet().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 假设训练到 epoch 10
for epoch in range(1, 11):
    # ... 训练代码 ...
    pass

# 保存 checkpoint
save_checkpoint(model, optimizer, epoch=10, loss=0.5, accuracy=0.89, 
                filepath='checkpoint_final.pt')

# 恢复训练
model_new = SimpleNet().to(device)
optimizer_new = optim.Adam(model_new.parameters(), lr=0.001)

checkpoint = load_checkpoint('checkpoint_final.pt', model_new, optimizer_new, device)
start_epoch = checkpoint['epoch'] + 1
print(f"从 epoch {start_epoch} 继续训练")

3.6 核心流程可视化

为了让你更直观地理解整个 checkpoint 的保存与恢复流程，我画了一张图：

Checkpoint 保存与恢复流程 训练循环 保存 checkpoint 模型 state_dict 优化器 state_dict 训练元数据 checkpoint.pt （pickle 序列化） 加载 checkpoint map_location 指定设备 恢复模型参数 恢复优化器状态 恢复训练进度 关键点：保存时打包所有状态，加载时按需恢复。优化器状态丢失会导致训练「冷启动」。

3.7 常见陷阱与避坑指南

最后，分享几个我踩过的坑：

  • 跨设备加载报错：我曾经在 GPU 上训练，想在 CPU 机器上做推理。直接 torch.load() 会报 CUDA 错误。解决方案就是加 map_location='cpu'。
  • 模型结构不匹配：保存的模型和加载时的模型结构必须完全一致。我试过改了网络层名字，结果 load_state_dict() 报 key 不匹配。建议用 strict=False 参数做部分加载，但要注意检查哪些 key 没匹配上。
  • 优化器参数组不一致：创建优化器时，如果参数组顺序或分组方式变了，加载优化器 state_dict 会失败。我的习惯是：保存和恢复时用完全相同的优化器初始化代码。
  • 大文件加载慢：几十 GB 的 checkpoint 加载可能很慢。我一般用 torch.load(..., mmap=True) 做内存映射加载，能显著减少内存占用。

我的习惯：每次保存 checkpoint 时，额外保存一个只包含模型参数的轻量版本（torch.save(model.state_dict(), 'model_only.pt')）。这样推理时只需要加载这个小文件，不用带优化器那些冗余数据。

好了，这一章的内容就到这里。state_dict 机制是 checkpoint 的基石，torch.save() 和 torch.load() 是操作工具。理解透这三者的关系，后续的分布式训练 checkpoint、断点续训等内容才能顺利展开。

← 上一章 📖 返回目录 下一章 →