PyTorch 模型的保存与加载全流程详解：第八篇 —— Save and Load the Model 教程

导语

深度学习模型训练通常耗时较长，只有掌握正确的“保存”和“恢复”方法，才能在实验、生产和迁移学习中高效复用你的模型成果。PyTorch 为我们提供了灵活易用的模型保存与加载机制，无论是仅保存参数权重还是连同训练状态与优化器一同存档，都能轻松胜任。

本章将带你系统掌握：

1. 1. PyTorch 的模型保存与加载机制原理；
2. 2. 只保存参数、保存完整训练状态的区别与场景；
3. 3. 保存与加载流程的详细实战代码；
4. 4. 跨设备（CPU/GPU）模型存档迁移方法；
5. 5. 多模型和多优化器同时保存、加载的进阶技巧；
6. 6. 常见错误与实用建议。

一、模型保存的原理与模式

概念讲解

PyTorch 保存模型有两种主流方式：

• o 保存参数（state_dict）：只保存参数权重，文件小、最常用，需代码中明确网络结构定义。
• o 保存整个模型（不推荐）：保存模型结构+参数，依赖 Python 代码环境、兼容性差。

最佳实践：一般仅保存 state_dict（即参数字典），而不是整个模型对象。这既保证了灵活性，也方便跨平台、跨版本迁移。

思考题

1. 1. 为什么官方建议优先保存 state_dict，而不是整个模型对象？
   答案：保存 state_dict 不依赖模型具体的 Python 类定义和代码环境，跨平台、跨版本兼容性好；直接保存模型对象则强依赖代码实现，可能因版本不一致等导致无法加载。

二、如何保存模型参数（state_dict）

概念讲解

• o state_dict 是一个简单的 Python 字典，存储所有参数和缓存（如卷积核、全连接层权重、偏置等）。
• o 使用 torch.save() 结合 model.state_dict()保存为 .pth 或 .pt 文件。

代码示例

import torch
from torch import nn

# 定义简单模型
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(28*28, 10)
    def forward(self, x):
        return self.fc(x)

model = NeuralNetwork()

# 保存参数字典
torch.save(model.state_dict(), "model_weights.pth")
print("模型参数已保存至 model_weights.pth")

思考题

1. 1. torch.save() 保存的本质是什么格式？
   答案：本质是通过 Python 的 pickle 协议将对象序列化保存为二进制文件。

三、如何加载模型参数

概念讲解

• o 加载时必须先重新定义与保存时完全一致的模型结构，然后使用 load_state_dict() 加载参数。
• o 加载后需切换为 eval() 模式（推理/验证），避免 Dropout、BatchNorm 等在测试阶段行为异常。

代码示例

# 1. 重新定义模型结构
model = NeuralNetwork()

# 2. 加载权重
model.load_state_dict(torch.load("model_weights.pth"))

# 3. 切换为评估模式
model.eval()
print("模型已成功加载并进入评估模式")

思考题

1. 1. 如果加载的参数与模型结构不一致会发生什么？
   答案：会报错，通常为参数维度不匹配，无法完成参数赋值。需保证模型结构与保存时完全一致。
2. 2. 训练时能否在加载后直接继续训练？
   答案：可以，加载参数后切换为 train() 模式即可继续训练。建议同时保存优化器等训练状态。

四、保存和恢复整个训练状态（进阶）

概念讲解

仅保存参数只能用于推理/迁移学习。如果要断点续训、完整恢复训练，建议保存“训练快照”：

• o 模型参数
• o 优化器状态（optimizer.state_dict()）
• o 学习率调度器状态（如有）
• o 训练进度指标（如 epoch、损失等）

代码示例：保存和加载完整训练状态

# 假设已有模型、优化器、调度器等
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
epoch = 10
loss = 0.05

# 保存完整快照
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'loss': loss
}, "checkpoint.pth")
print("完整训练状态已保存至 checkpoint.pth")

# 恢复快照
checkpoint = torch.load("checkpoint.pth")
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

print(f"已恢复到第 {epoch} 轮，损失为 {loss}")

思考题

1. 1. 为什么断点训练时必须同时恢复优化器状态？
   答案：优化器包含动量、自适应学习率等信息，若只恢复模型参数而忽略优化器，可能导致模型训练轨迹异常、不连贯。

五、跨设备保存与加载（CPU/GPU）

概念讲解

• o 保存时模型在 CPU 或 GPU 无影响，加载时需指定目标设备。
• o 常用场景：用 GPU 训练后需在 CPU 上推理/测试，或在不同硬件间迁移。

代码示例

# 保存时（无关设备）
torch.save(model.state_dict(), "model_gpu.pth")

# 加载到 CPU
model_cpu = NeuralNetwork()
model_cpu.load_state_dict(torch.load("model_gpu.pth", map_location='cpu'))

# 加载到 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_gpu = NeuralNetwork().to(device)
model_gpu.load_state_dict(torch.load("model_gpu.pth", map_location=device))

思考题

1. 1. 如果保存的是 GPU 版本参数，直接在 CPU 上加载会怎样？
   答案：只要用 map_location='cpu'，PyTorch 会自动转换参数为 CPU 张量，无需手动处理。

六、多模型、多优化器联合保存与恢复

概念讲解

实际项目（如 GAN、对抗训练、多任务）可能涉及多个模型和优化器。建议统一以字典方式组织，方便扩展和管理。

代码示例

# 假设有生成器和判别器
generator = NeuralNetwork()
discriminator = NeuralNetwork()
opt_g = torch.optim.Adam(generator.parameters())
opt_d = torch.optim.Adam(discriminator.parameters())

# 联合保存
torch.save({
    'generator_state_dict': generator.state_dict(),
    'discriminator_state_dict': discriminator.state_dict(),
    'opt_g_state_dict': opt_g.state_dict(),
    'opt_d_state_dict': opt_d.state_dict()
}, "gan_checkpoint.pth")

# 恢复
checkpoint = torch.load("gan_checkpoint.pth")
generator.load_state_dict(checkpoint['generator_state_dict'])
discriminator.load_state_dict(checkpoint['discriminator_state_dict'])
opt_g.load_state_dict(checkpoint['opt_g_state_dict'])
opt_d.load_state_dict(checkpoint['opt_d_state_dict'])

思考题

1. 1. 如果只想恢复一部分模型的参数，可以怎么做？
   答案：只从 checkpoint 中取出并加载所需部分即可，其他模型参数不受影响。

七、模型保存与加载的常见错误与建议

• o 结构不一致导致加载失败：始终保证模型结构与保存时一致（包括参数名、层数、形状）。
• o 保存/加载时 device 不一致：用 map_location 显式指定目标设备，避免 device 不匹配报错。
• o 保存模型对象而非 state_dict：强依赖代码实现和环境，易导致兼容性问题，不建议直接 torch.save(model, ...)。
• o 未保存优化器等训练状态：断点训练务必保存 optimizer、scheduler 等状态。

调试技巧

# 检查所有参数名和形状
state_dict = torch.load("model_weights.pth")
for k, v in state_dict.items():
    print(k, v.shape)

八、生产环境部署建议

• o 模型导出 ONNX：用于跨平台（如 TensorRT、OpenVINO）推理，可用 torch.onnx.export。
• o 保存配置文件：记录模型结构超参数、预处理流程、训练环境等，保证重现实验。
• o 版本管理与备份：合理命名保存文件，定期备份，方便实验管理和回溯。

知识小结

• o PyTorch 推荐以 state_dict 方式保存和加载模型，灵活、兼容性好；
• o 断点训练和迁移学习需完整保存优化器、调度器等训练状态；
• o 加载时注意模型结构与设备匹配，可用 map_location 实现跨 CPU/GPU 迁移；
• o 多模型、多优化器统一保存字典，便于复杂场景扩展；
• o 部署与生产环境可结合 ONNX 导出等方式，提升可移植性。