【深度学习-Day 29】PyTorch模型持久化指南：从保存到部署的第一步

Langchain系列文章目录

01-玩转LangChain：从模型调用到Prompt模板与输出解析的完整指南
 02-玩转 LangChain Memory 模块：四种记忆类型详解及应用场景全覆盖
 03-全面掌握 LangChain：从核心链条构建到动态任务分配的实战指南
 04-玩转 LangChain：从文档加载到高效问答系统构建的全程实战
 05-玩转 LangChain：深度评估问答系统的三种高效方法（示例生成、手动评估与LLM辅助评估）
06-从 0 到 1 掌握 LangChain Agents：自定义工具 + LLM 打造智能工作流！
07-【深度解析】从GPT-1到GPT-4：ChatGPT背后的核心原理全揭秘
 08-【万字长文】MCP深度解析：打通AI与世界的"USB-C"，模型上下文协议原理、实践与未来

Python系列文章目录

PyTorch系列文章目录

机器学习系列文章目录

深度学习系列文章目录

Java系列文章目录

JavaScript系列文章目录

深度学习系列文章目录

29-【深度学习-Day 29】PyTorch模型持久化指南：从保存到部署的第一步

文章目录

Langchain系列文章目录
Python系列文章目录
PyTorch系列文章目录
机器学习系列文章目录
深度学习系列文章目录
Java系列文章目录
JavaScript系列文章目录
深度学习系列文章目录
前言
一、为何要保存模型？场景驱动的需求
- [1.1 推理与部署 (Inference & Deployment)](#1.1 推理与部署 (Inference & Deployment))
- [1.2 断点续训 (Checkpointing)](#1.2 断点续训 (Checkpointing))
- [1.3 模型分享与复现 (Sharing & Reproducibility)](#1.3 模型分享与复现 (Sharing & Reproducibility))
二、保存策略：整存还是零取？
- [2.1 只保存模型参数 (State Dictionary)](#2.1 只保存模型参数 (State Dictionary))
- - [2.1.1 什么是 `state_dict`？](#2.1.1 什么是 state_dict？)
  - [2.1.2 优点与适用场景](#2.1.2 优点与适用场景)
- [2.2 保存完整模型 (Entire Model)](#2.2 保存完整模型 (Entire Model))
- - [2.2.1 工作原理](#2.2.1 工作原理)
  - [2.2.2 优点与潜在问题](#2.2.2 优点与潜在问题)
- [2.3 策略对比与选择](#2.3 策略对比与选择)
[三、PyTorch 实战：模型保存与加载](#三、PyTorch 实战：模型保存与加载)
- [3.1 准备工作：定义并训练一个简单模型](#3.1 准备工作：定义并训练一个简单模型)
- [3.2 实践一：保存和加载模型参数 (`state_dict`)](#3.2 实践一：保存和加载模型参数 (state_dict))
- - - [（1）保存 `state_dict`](#（1）保存 state_dict)
    - [（2）加载 `state_dict`](#（2）加载 state_dict)
- [3.3 实践二：保存和加载完整模型](#3.3 实践二：保存和加载完整模型)
- - - [（1）保存完整模型](#（1）保存完整模型)
    - [（2）加载完整模型](#（2）加载完整模型)
- [3.4 进阶：保存训练检查点 (Checkpointing)](#3.4 进阶：保存训练检查点 (Checkpointing))
- - - [（1）保存检查点](#（1）保存检查点)
    - [（2）加载检查点以恢复训练](#（2）加载检查点以恢复训练)
四、常见问题与最佳实践
- [4.1 `.pt`, `.pth`, `.pkl`？用什么扩展名？](#4.1 .pt, .pth, .pkl？用什么扩展名？)
- [4.2 加载模型到不同设备 (CPU/GPU)](#4.2 加载模型到不同设备 (CPU/GPU))
- [4.3 最佳实践总结](#4.3 最佳实践总结)
五、总结

前言

恭喜你，通过前面一系列的学习和实践，你已经能够成功训练出一个深度学习模型了！但是，当Jupyter Notebook关闭，或者训练过程意外中断，那些花费了数小时甚至数天计算资源得到的模型参数难道就要付诸东流吗？当然不。模型训练的最终目的，是为了将其应用于实际场景，或者在未来继续优化。这就引出了深度学习流程中至关重要的一环------模型保存与加载。

本篇文章将带你深入理解为什么需要保存模型，探索不同的保存策略，并手把手教你如何使用 PyTorch 实现模型的持久化、加载和复用。掌握这项技能，意味着你的模型不再是一次性的"消耗品"，而是可以随时待命、持续进化的宝贵资产。

一、为何要保存模型？场景驱动的需求

在深入技术细节之前，我们首先要明确保存模型的动机。它绝不仅仅是"备份"那么简单，而是贯穿模型整个生命周期的核心需求。

1.1 推理与部署 (Inference & Deployment)

这是最常见的需求。模型训练完成后，我们需要将其部署到生产环境中，例如一个网站后端、一个移动App或者一个边缘设备上，来对新的、未见过的数据进行预测（这个过程称为"推理"）。显然，我们不可能在每个需要预测的地方都重新训练一遍模型。因此，必须将训练好的模型状态保存下来，在推理环境中直接加载使用。

1.2 断点续训 (Checkpointing)

深度学习模型的训练过程往往非常耗时，从几小时到几周不等。如果训练中途因为断电、程序崩溃等原因而中断，没有保存机制的话，一切都得从头再来，这将是巨大的时间和资源浪费。**检查点（Checkpointing）**机制允许我们在训练过程中定期保存模型的状态（包括权重、训练到第几轮、优化器状态等），以便在中断后能够从上次保存的地方无缝地继续训练。

学术研究和团队协作中，我们需要将自己的模型分享给他人，以便他们能够复现我们的实验结果。一个标准的模型保存文件是保证研究可复现性的关键。此外，在迁移学习中，我们常常加载他人预训练好的模型，并在其基础上进行微调，这也依赖于成熟的模型保存与加载体系。

二、保存策略：整存还是零取？

在 PyTorch 中，保存模型主要有两种方式：只保存模型参数 和保存完整模型。这两种方式各有优劣，理解它们的区别对于选择合适的策略至关重要。

2.1 只保存模型参数 (State Dictionary)

这是一种更推荐、更灵活、也更"Pythonic"的方式。它只保存模型的"状态字典"(state_dict)。

2.1.1 什么是 `state_dict`？

state_dict 是一个 Python 字典对象，它将模型的每一个可学习的参数层（如卷积层、全连接层）映射到其对应的参数张量（权重 weight 和偏置 bias）。

举个例子，一个简单的模型，其 state_dict 可能长这样：

python 复制代码

{
    'conv1.weight': tensor([...]),
    'conv1.bias': tensor([...]),
    'fc1.weight': tensor([...]),
    'fc1.bias': tensor([...])
}

它就像是模型的"骨架蓝图"，只包含了最核心的参数信息，不包含模型的结构定义代码。

2.1.2 优点与适用场景

灵活性与可移植性高：由于只保存参数，模型结构代码和参数文件是分离的。这意味着你可以轻松地将这些参数加载到一个结构相同但代码实现略有不同的模型中。当你的项目代码重构，或者PyTorch版本更新时，这种方式的兼容性最好。
文件更小：通常比保存整个模型的文件要小。

2.2 保存完整模型 (Entire Model)

这种方式使用 Python 的 pickle 模块将整个模型对象（包括其结构、参数等所有信息）序列化到磁盘。

2.2.1 工作原理

当你保存整个模型时，PyTorch 不仅保存了 state_dict，还保存了定义模型类的代码路径和结构。加载时，它会尝试按照保存的路径去寻找并重建这个模型对象。

2.2.2 优点与潜在问题

优点：语法简单，保存和加载都只需要一行代码，非常直观。
潜在问题：这种方式将模型代码与数据紧紧地"捆绑"在了一起。如果加载模型的项目目录结构发生变化，或者你重命名了包含模型定义的文件，加载时就可能找不到对应的类而报错。此外，跨 PyTorch 版本的兼容性也较差。

2.3 策略对比与选择

为了更直观地比较，我们用一个表格来总结：

特性	只保存参数 (`state_dict`)	保存完整模型
推荐度	⭐⭐⭐⭐⭐ (推荐)	⭐⭐
灵活性	高，代码与数据分离	低，代码与数据耦合
可移植性	高，易于代码重构和跨版本	低，对代码路径和版本敏感
安全性	较高	较低（`pickle`可能存在安全风险）
文件大小	较小	较大
使用便利性	加载时需先创建模型实例	非常简单

三、PyTorch 实战：模型保存与加载

理论讲完了，让我们卷起袖子，用代码来实践一下。我们将使用一个简单的多层感知机（MLP）模型作为示例。

3.1 准备工作：定义并训练一个简单模型

首先，我们定义一个用于图像分类的简单网络，并假设我们已经有了数据加载器和训练循环（为聚焦主题，此处省略训练细节）。

python 复制代码

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

# 1. 定义模型结构
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

# 2. 初始化模型、损失函数和优化器
model = SimpleNet()
# 假设模型在GPU上训练
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

optimizer = optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()

# 3. (伪代码) 假设我们已经训练了一段时间...
print("模型训练完成或到达一个检查点...")
# 假设训练了 5 个 epoch
epoch = 5
current_loss = 0.5 # 假设这是最新的loss

3.2 实践一：保存和加载模型参数 (`state_dict`)

这是最推荐的方式。

（1）保存 `state_dict`

python 复制代码

# 定义保存路径
PATH_STATE_DICT = "model_state_dict.pth"

# 保存模型的状态字典
print(f"正在保存模型参数到 {PATH_STATE_DICT}...")
torch.save(model.state_dict(), PATH_STATE_DICT)
print("保存完成!")

（2）加载 `state_dict`

要加载参数，我们必须先创建一个与被保存时结构完全相同的模型实例。

python 复制代码

# 1. 创建一个新的模型实例
loaded_model_from_dict = SimpleNet()

# 2. 加载状态字典
print(f"正在从 {PATH_STATE_DICT} 加载模型参数...")
loaded_model_from_dict.load_state_dict(torch.load(PATH_STATE_DICT))
print("加载完成!")

# 3. 将模型设置为评估模式 (非常重要!)
# 这会关闭 Dropout 和 BatchNorm 等层的训练行为。
loaded_model_from_dict.eval()

# 4. (可选) 将模型移动到合适的设备
loaded_model_from_dict.to(device)

# 现在，loaded_model_from_dict 可以用于推理了
# with torch.no_grad():
#     new_data = ...
#     prediction = loaded_model_from_dict(new_data)

关键点 ：加载后一定要调用 .eval() 方法，这是新手常犯的错误之一。

3.3 实践二：保存和加载完整模型

虽然不推荐，但了解其用法也很有必要。

（1）保存完整模型

python 复制代码

# 定义保存路径
PATH_FULL_MODEL = "full_model.pt"

# 保存整个模型
print(f"正在保存完整模型到 {PATH_FULL_MODEL}...")
torch.save(model, PATH_FULL_MODEL)
print("保存完成!")

（2）加载完整模型

加载时，我们不需要预先创建模型实例。

python 复制代码

# 直接加载整个模型对象
print(f"正在从 {PATH_FULL_MODEL} 加载完整模型...")
loaded_full_model = torch.load(PATH_FULL_MODEL)
print("加载完成!")

# 同样，设置为评估模式
loaded_full_model.eval()

# (可选) 移动到设备
loaded_full_model.to(device)

# 模型已准备好用于推理

3.4 进阶：保存训练检查点 (Checkpointing)

对于断点续训，我们不仅需要模型参数，还需要优化器状态、当前的 epoch 数、损失等信息。最佳实践是创建一个字典来统一保存它们。

（1）保存检查点

python 复制代码

CHECKPOINT_PATH = "training_checkpoint.pt"

print(f"正在保存训练检查点到 {CHECKPOINT_PATH}...")
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': current_loss,
    # ... 还可以保存其他任何你需要的信息
}, CHECKPOINT_PATH)
print("检查点保存完毕!")

（2）加载检查点以恢复训练

python 复制代码

# 1. 像往常一样，先初始化模型和优化器
resume_model = SimpleNet()
resume_optimizer = optim.SGD(resume_model.parameters(), lr=0.01) # lr可以先随意设置

# 2. 加载检查点字典
checkpoint = torch.load(CHECKPOINT_PATH)

# 3. 分别恢复模型和优化器的状态
resume_model.load_state_dict(checkpoint['model_state_dict'])
resume_optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch']
last_loss = checkpoint['loss']

# 4. 别忘了将模型设置为训练模式
resume_model.train() 

print(f"成功从 epoch {start_epoch} 恢复训练。上一轮损失: {last_loss:.4f}")
# 现在你可以从 start_epoch + 1 开始继续你的训练循环了
# for epoch in range(start_epoch + 1, num_epochs):
#     # ... training loop ...

四、常见问题与最佳实践

4.1 `.pt`, `.pth`, `.pkl`？用什么扩展名？

PyTorch 保存模型时，对文件扩展名没有强制要求。.pt、.pth 是社区中常见的约定，torch.save 内部使用的是 Python 的 pickle，所以 .pkl 也能工作。
建议：使用 .pt 或 .pth，这能清晰地表明文件内容是 PyTorch 相关的数据。

4.2 加载模型到不同设备 (CPU/GPU)

一个常见的场景是在带 GPU 的服务器上训练和保存模型，然后在只有 CPU 的本地机器或服务器上加载它用于推理。如果直接加载，PyTorch会报错，因为它会尝试将模型参数加载到原始的 GPU 设备上。

解决方案 ：在 torch.load() 中使用 map_location 参数。

python 复制代码

# 假设模型在 GPU 上保存，我们想在 CPU 上加载
# 方法一：加载到 CPU
cpu_device = torch.device('cpu')
model_state = torch.load(PATH_STATE_DICT, map_location=cpu_device)
model.load_state_dict(model_state)

# 方法二：更通用的写法
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_state = torch.load(PATH_STATE_DICT, map_location=device)
model.load_state_dict(model_state)
model.to(device)

map_location 会告诉 torch.load 将所有张量加载到指定的设备上，从而避免设备不匹配的错误。

4.3 最佳实践总结

优先使用 state_dict：为了代码的健壮性和可移植性，总是优先选择只保存和加载模型参数。
保存检查点以续训：对于长时间的训练任务，务必实现检查点机制，将模型、优化器状态和训练轮次等信息打包保存在一个字典中。
模式切换要牢记 ：加载模型用于推理前，调用 .eval()；加载检查点恢复训练前，调用 .train()。
考虑设备差异 ：在加载模型时，使用 map_location 参数来优雅地处理 CPU/GPU 的差异。
版本信息 ：在保存检查点时，可以考虑一并保存 PyTorch 的版本号 (torch.__version__)，这对于未来的调试和复现非常有帮助。

五、总结

掌握模型的保存与加载，是连接深度学习理论研究与实际应用的关键桥梁。它让我们的辛勤训练的成果得以固化、复用和分享。通过本文的学习，我们应掌握以下核心要点：

重要性 ：模型保存是实现推理部署 、断点续训 和分享复现这三大核心应用场景的基础。
核心策略 ：PyTorch 提供两种保存方式。只保存参数 (state_dict) 是官方推荐的最佳实践，因为它灵活、健壮且易于维护。保存完整模型虽然简单，但因其脆弱性，应避免在生产和协作环境中使用。
实战技能 ：我们学习了如何通过 torch.save() 和 model.load_state_dict() 来完成标准模型的保存与加载。更重要的是，我们掌握了如何构建和恢复训练检查点（Checkpoint），以实现稳健的断点续训。
关键细节 ：牢记在加载模型后根据用途调用 .eval()（推理）或 .train()（续训）模式，并善用 map_location 处理设备差异，这些是保证代码正确运行的"护身符"。

现在，你不仅能训练模型，更能驾驭它们，让它们走出实验室，准备好在更广阔的世界中创造价值。这是你从"炼丹师"向"算法工程师"迈出的坚实一步。

【深度学习-Day 29】PyTorch模型持久化指南：从保存到部署的第一步

Langchain系列文章目录

Python系列文章目录

PyTorch系列文章目录

机器学习系列文章目录

深度学习系列文章目录

Java系列文章目录

JavaScript系列文章目录

深度学习系列文章目录

文章目录

前言

一、为何要保存模型？场景驱动的需求

1.1 推理与部署 (Inference & Deployment)

1.2 断点续训 (Checkpointing)

1.3 模型分享与复现 (Sharing & Reproducibility)

二、保存策略：整存还是零取？

2.1 只保存模型参数 (State Dictionary)

2.1.1 什么是 state_dict？

2.1.2 优点与适用场景

2.2 保存完整模型 (Entire Model)

2.2.1 工作原理

2.2.2 优点与潜在问题

2.3 策略对比与选择

三、PyTorch 实战：模型保存与加载

3.1 准备工作：定义并训练一个简单模型

3.2 实践一：保存和加载模型参数 (state_dict)

（1） 保存 state_dict

（2） 加载 state_dict

3.3 实践二：保存和加载完整模型

（1） 保存完整模型

（2） 加载完整模型

3.4 进阶：保存训练检查点 (Checkpointing)

（1） 保存检查点

（2） 加载检查点以恢复训练

四、常见问题与最佳实践

4.1 .pt, .pth, .pkl？用什么扩展名？

4.2 加载模型到不同设备 (CPU/GPU)

4.3 最佳实践总结

五、总结

2.1.1 什么是 `state_dict`？

3.2 实践一：保存和加载模型参数 (`state_dict`)

（1）保存 `state_dict`

（2）加载 `state_dict`

（1）保存完整模型

（2）加载完整模型

（1）保存检查点

（2）加载检查点以恢复训练

4.1 `.pt`, `.pth`, `.pkl`？用什么扩展名？