深度学习模型保存与装载：从训练到部署的全流程指南

最新推荐文章于 2025-09-08 20:47:15 发布

二向箔reverse

最新推荐文章于 2025-09-08 20:47:15 发布

阅读量1k

点赞数 37

CC 4.0 BY-SA版权

文章标签：深度学习 pytorch 分类

本文链接：https://blog.csdn.net/cpluspluszz/article/details/151122439

在深度学习项目中，模型训练往往需要数小时、数天甚至数周的时间。一旦训练中断（如硬件故障、代码错误），重新开始训练会浪费大量时间；更关键的是，当模型达到理想效果时，我们需要将其保存以便后续推理、部署或进一步优化。本文将围绕PyTorch框架，详细讲解模型保存与装载的核心技巧，结合实际代码场景，帮你掌握这一深度学习必备技能。

一、为什么需要保存与装载模型？

假设你正在训练一个图像分类模型，经过24小时的训练，验证集准确率达到了92%，但此时电脑突然断电——如果没有保存模型，所有努力将付诸东流。模型保存与装载的核心价值在于：

断点续训：保存训练过程中的中间状态（如模型参数、优化器状态、当前epoch），避免因意外中断导致的重复训练；
模型部署：将训练好的模型保存为文件，用于生产环境的实时推理（如在线图片分类API）；
迁移学习：加载预训练模型（如ResNet）的特征提取层，仅微调全连接层以适应新任务；
实验对比：保存不同超参数下的模型版本，方便后续对比实验效果。

二、PyTorch模型保存的核心：`state_dict`与完整模型

在PyTorch中，模型保存的核心是state_dict（状态字典）——一个包含模型所有可学习参数（如卷积层权重、全连接层偏置）的Python字典。此外，还可以保存优化器状态、当前训练轮次（epoch）、损失值等信息，实现完整的训练状态保存。

2.1 保存`state_dict`（推荐方式）

state_dict仅保存模型的参数，不包含模型结构信息。这种方式更灵活，适用于：

模型结构可能修改的场景（如调整层数后加载旧参数）；
跨平台部署（只需保存参数，部署时重新定义模型结构）。

保存代码示例：

# 在训练循环中保存最佳模型（基于用户提供的训练代码扩展）
best_acc = 0.0
save_path = "best_model.pth"  # 保存路径

for epoch in range(100):
    train(train_dataloader, model, loss_df, optimizer)
    current_acc = test(test_dataloader, model)  # 假设test函数返回当前准确率
    
    # 如果当前准确率优于历史最佳，保存模型
    if current_acc > best_acc:
        best_acc = current_acc
        # 保存state_dict（模型参数）
        torch.save(model.state_dict(), save_path)
        print(f"Epoch {epoch+1}: 验证准确率提升至{current_acc:.4f}，模型已保存至{save_path}")

2.2 保存完整模型（含结构）

若需要将模型结构与参数一起保存（如快速部署时无需重新定义模型），可以直接保存整个模型对象。但这种方式依赖模型类的定义，跨代码版本或不同机器时可能因类定义不一致导致加载失败。

保存代码示例：

# 保存完整模型（结构+参数）
full_model_path = "full_model.pth"
torch.save(model, full_model_path)  # 直接保存模型对象

2.3 保存训练状态（含优化器、epoch等）

为了实现断点续训（从上次中断的位置继续训练），需要保存更多上下文信息：模型参数、优化器状态、当前epoch、当前损失值等。

保存代码示例：

checkpoint_path = "checkpoint.pth"
# 保存训练状态
torch.save({
    'epoch': epoch,                  # 当前训练轮次
    'model_state_dict': model.state_dict(),  # 模型参数
    'optimizer_state_dict': optimizer.state_dict(),  # 优化器状态（如SGD的动量参数）
    'loss': loss,                    # 当前批次损失值
}, checkpoint_path)

三、模型装载：从文件到推理/训练

加载模型时，需根据保存的内容（state_dict、完整模型或训练状态）选择对应的加载方式，并注意设备兼容性（如从GPU保存的模型加载到CPU）。

3.1 加载`state_dict`（最常用）

加载state_dict时，需要先初始化模型结构，再通过load_state_dict()方法加载参数。这种方式灵活，适用于：

加载预训练模型（如ImageNet预训练的ResNet）；
恢复训练（结合优化器状态）；
跨设备加载（如从GPU保存的模型加载到CPU）。

加载代码示例：

# 初始化模型结构（需与保存时的模型类完全一致）
model = CNN().to(device)  # device是'cuda'或'cpu'

# 加载保存的state_dict
saved_state = torch.load(save_path)  # 加载参数文件
model.load_state_dict(saved_state)   # 将参数加载到模型中

# 验证加载是否成功（可选）
test(test_dataloader, model)  # 输出准确率，确认模型有效

跨设备加载（GPU→CPU）

若模型在GPU上训练，现在需要在CPU上推理，加载时需指定map_location参数：

# 从GPU保存的模型加载到CPU
model = CNN()
saved_state = torch.load(save_path, map_location=torch.device('cpu'))  # 强制加载到CPU
model.load_state_dict(saved_state)

加载部分参数（迁移学习）

当需要复用预训练模型的部分层时，可以只加载匹配的参数：

# 假设预训练模型有conv1层，当前模型也有conv1层
pretrained_state = torch.load("pretrained_model.pth")
model_state = model.state_dict()

# 筛选出匹配的参数（键名相同且形状一致）
matched_params = {k: v for k, v in pretrained_state.items() if k in model_state and model_state[k].shape == v.shape}

# 加载匹配的参数
model_state.update(matched_params)
model.load_state_dict(model_state)

3.2 加载完整模型

若保存的是完整模型（含结构），加载时无需重新定义模型类，直接调用torch.load()即可。但需注意：模型类必须在当前代码环境中定义，否则会报错。

加载代码示例：

# 直接加载完整模型（需提前定义CNN类）
model = torch.load(full_model_path)
model.to(device)  # 加载到目标设备（GPU/CPU）

# 推理示例
img = Image.open("test_image.jpg")
transformed_img = data_transformsimg.unsqueeze(0).to(device)  # 预处理并增加batch维度
with torch.no_grad():  # 关闭梯度计算（推理时无需）
    pred = model(transformed_img)
print(f"预测类别：{pred.argmax(1).item()}")

3.3 恢复训练（断点续训）

若需要从上次中断的位置继续训练，需同时加载模型参数、优化器状态和训练进度（如当前epoch）。

恢复训练代码示例：

# 初始化模型和优化器
model = CNN().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 加载检查点（包含训练状态）
checkpoint = torch.load(checkpoint_path)
start_epoch = checkpoint['epoch'] + 1  # 从下一轮开始
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
last_loss = checkpoint['loss']

print(f"恢复训练：从第{start_epoch}轮开始，上一轮损失为{last_loss:.4f}")

# 继续训练循环
for epoch in range(start_epoch, 100):
    train(train_dataloader, model, loss_df, optimizer)
    current_acc = test(test_dataloader, model)
    # ...（保存逻辑同上）

四、模型保存的最佳实践

4.1 命名规范

按实验标识命名：如resnet34_lr0.01_epoch90_acc92.pth，包含模型类型、超参数、关键指标，方便后续查找；
保存最佳模型和最终模型：best_val_acc.pth（验证集最佳）和final_model.pth（训练结束时的模型）。

4.2 定期保存

在长耗时训练中，建议每N个epoch保存一次模型（如每10轮），避免因意外中断丢失过多进度：

if epoch % 10 == 0:  # 每10轮保存一次
    torch.save(model.state_dict(), f"model_epoch{epoch}.pth")

4.3 验证保存的模型

保存后建议加载模型并验证其性能（如在验证集上测试准确率），确保保存过程无错误。

五、常见问题与解决方案

问题1：加载模型时报错“Missing key(s) in state_dict”

原因：当前模型结构与保存时的模型结构不一致（如修改了层数或参数名）。
解决：

检查模型类定义是否与保存时完全一致；
若需加载部分参数，使用strict=False参数（忽略不匹配的键）：
```
model.load_state_dict(torch.load(save_path), strict=False)
```

问题2：加载GPU模型到CPU时报错“CUDA error: an illegal memory access was encountered”

原因：模型参数存储在GPU内存中，但当前环境无GPU或未正确指定设备。
解决：加载时通过map_location指定目标设备：

model.load_state_dict(torch.load(save_path, map_location=torch.device('cpu')))

问题3：优化器状态加载后训练效果异常

原因：优化器状态（如SGD的动量缓存）与当前模型参数不匹配（如学习率被修改）。
解决：

确保加载优化器状态时，学习率等超参数与保存时一致；

若需调整学习率，可在加载后手动修改优化器的param_groups：

optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
for param_group in optimizer.param_groups:
    param_group['lr'] = 0.001  # 调整学习率

六、总结

模型保存与装载是深度学习工作流中不可或缺的一环，掌握这一技能能显著提升开发效率。本文重点讲解了：

保存策略：state_dict（灵活）、完整模型（便捷）、训练状态（断点续训）；
加载场景：推理、继续训练、迁移学习；
最佳实践：命名规范、定期保存、验证模型；
常见问题：结构不匹配、设备不兼容、优化器状态异常。

回到用户提供的代码，在训练循环中添加模型保存逻辑（如保存最佳验证模型），并在需要时加载模型进行推理或继续训练，即可轻松实现模型的持久化管理。下次训练时，再也不用担心中断或重复劳动啦！# 深度学习模型保存与装载：从训练到部署的全流程指南

一、为什么需要保存与装载模型？

断点续训：保存训练过程中的中间状态（如模型参数、优化器状态、当前epoch），避免因意外中断导致的重复训练；
模型部署：将训练好的模型保存为文件，用于生产环境的实时推理（如在线图片分类API）；
迁移学习：加载预训练模型（如ResNet）的特征提取层，仅微调全连接层以适应新任务；
实验对比：保存不同超参数下的模型版本，方便后续对比实验效果。

二、PyTorch模型保存的核心：`state_dict`与完整模型

2.1 保存`state_dict`（推荐方式）

state_dict仅保存模型的参数，不包含模型结构信息。这种方式更灵活，适用于：

模型结构可能修改的场景（如调整层数后加载旧参数）；
跨平台部署（只需保存参数，部署时重新定义模型结构）。

保存代码示例：

# 在训练循环中保存最佳模型（基于用户提供的训练代码扩展）
best_acc = 0.0
save_path = "best_model.pth"  # 保存路径

for epoch in range(100):
    train(train_dataloader, model, loss_df, optimizer)
    current_acc = test(test_dataloader, model)  # 假设test函数返回当前准确率
    
    # 如果当前准确率优于历史最佳，保存模型
    if current_acc > best_acc:
        best_acc = current_acc
        # 保存state_dict（模型参数）
        torch.save(model.state_dict(), save_path)
        print(f"Epoch {epoch+1}: 验证准确率提升至{current_acc:.4f}，模型已保存至{save_path}")

2.2 保存完整模型（含结构）

保存代码示例：

# 保存完整模型（结构+参数）
full_model_path = "full_model.pth"
torch.save(model, full_model_path)  # 直接保存模型对象

2.3 保存训练状态（含优化器、epoch等）

为了实现断点续训（从上次中断的位置继续训练），需要保存更多上下文信息：模型参数、优化器状态、当前epoch、当前损失值等。

保存代码示例：

checkpoint_path = "checkpoint.pth"
# 保存训练状态
torch.save({
    'epoch': epoch,                  # 当前训练轮次
    'model_state_dict': model.state_dict(),  # 模型参数
    'optimizer_state_dict': optimizer.state_dict(),  # 优化器状态（如SGD的动量参数）
    'loss': loss,                    # 当前批次损失值
}, checkpoint_path)

三、模型装载：从文件到推理/训练

加载模型时，需根据保存的内容（state_dict、完整模型或训练状态）选择对应的加载方式，并注意设备兼容性（如从GPU保存的模型加载到CPU）。

3.1 加载`state_dict`（最常用）

加载state_dict时，需要先初始化模型结构，再通过load_state_dict()方法加载参数。这种方式灵活，适用于：

加载预训练模型（如ImageNet预训练的ResNet）；
恢复训练（结合优化器状态）；
跨设备加载（如从GPU保存的模型加载到CPU）。

加载代码示例：

# 初始化模型结构（需与保存时的模型类完全一致）
model = CNN().to(device)  # device是'cuda'或'cpu'

# 加载保存的state_dict
saved_state = torch.load(save_path)  # 加载参数文件
model.load_state_dict(saved_state)   # 将参数加载到模型中

# 验证加载是否成功（可选）
test(test_dataloader, model)  # 输出准确率，确认模型有效

跨设备加载（GPU→CPU）

若模型在GPU上训练，现在需要在CPU上推理，加载时需指定map_location参数：

# 从GPU保存的模型加载到CPU
model = CNN()
saved_state = torch.load(save_path, map_location=torch.device('cpu'))  # 强制加载到CPU
model.load_state_dict(saved_state)

加载部分参数（迁移学习）

当需要复用预训练模型的部分层时，可以只加载匹配的参数：

# 假设预训练模型有conv1层，当前模型也有conv1层
pretrained_state = torch.load("pretrained_model.pth")
model_state = model.state_dict()

# 筛选出匹配的参数（键名相同且形状一致）
matched_params = {k: v for k, v in pretrained_state.items() if k in model_state and model_state[k].shape == v.shape}

# 加载匹配的参数
model_state.update(matched_params)
model.load_state_dict(model_state)

3.2 加载完整模型

加载代码示例：

# 直接加载完整模型（需提前定义CNN类）
model = torch.load(full_model_path)
model.to(device)  # 加载到目标设备（GPU/CPU）

# 推理示例
img = Image.open("test_image.jpg")
transformed_img = data_transformsimg.unsqueeze(0).to(device)  # 预处理并增加batch维度
with torch.no_grad():  # 关闭梯度计算（推理时无需）
    pred = model(transformed_img)
print(f"预测类别：{pred.argmax(1).item()}")

3.3 恢复训练（断点续训）

若需要从上次中断的位置继续训练，需同时加载模型参数、优化器状态和训练进度（如当前epoch）。

恢复训练代码示例：

# 初始化模型和优化器
model = CNN().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 加载检查点（包含训练状态）
checkpoint = torch.load(checkpoint_path)
start_epoch = checkpoint['epoch'] + 1  # 从下一轮开始
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
last_loss = checkpoint['loss']

print(f"恢复训练：从第{start_epoch}轮开始，上一轮损失为{last_loss:.4f}")

# 继续训练循环
for epoch in range(start_epoch, 100):
    train(train_dataloader, model, loss_df, optimizer)
    current_acc = test(test_dataloader, model)
    # ...（保存逻辑同上）

四、模型保存的最佳实践

4.1 命名规范

按实验标识命名：如resnet34_lr0.01_epoch90_acc92.pth，包含模型类型、超参数、关键指标，方便后续查找；
保存最佳模型和最终模型：best_val_acc.pth（验证集最佳）和final_model.pth（训练结束时的模型）。

4.2 定期保存

在长耗时训练中，建议每N个epoch保存一次模型（如每10轮），避免因意外中断丢失过多进度：

if epoch % 10 == 0:  # 每10轮保存一次
    torch.save(model.state_dict(), f"model_epoch{epoch}.pth")

4.3 验证保存的模型

保存后建议加载模型并验证其性能（如在验证集上测试准确率），确保保存过程无错误。

五、常见问题与解决方案

问题1：加载模型时报错“Missing key(s) in state_dict”

原因：当前模型结构与保存时的模型结构不一致（如修改了层数或参数名）。
解决：

检查模型类定义是否与保存时完全一致；
若需加载部分参数，使用strict=False参数（忽略不匹配的键）：
```
model.load_state_dict(torch.load(save_path), strict=False)
```

问题2：加载GPU模型到CPU时报错“CUDA error: an illegal memory access was encountered”

原因：模型参数存储在GPU内存中，但当前环境无GPU或未正确指定设备。
解决：加载时通过map_location指定目标设备：

model.load_state_dict(torch.load(save_path, map_location=torch.device('cpu')))

问题3：优化器状态加载后训练效果异常

原因：优化器状态（如SGD的动量缓存）与当前模型参数不匹配（如学习率被修改）。
解决：

确保加载优化器状态时，学习率等超参数与保存时一致；

若需调整学习率，可在加载后手动修改优化器的param_groups：

optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
for param_group in optimizer.param_groups:
    param_group['lr'] = 0.001  # 调整学习率

六、总结

模型保存与装载是深度学习工作流中不可或缺的一环，掌握这一技能能显著提升开发效率。本文重点讲解了：

保存策略：state_dict（灵活）、完整模型（便捷）、训练状态（断点续训）；
加载场景：推理、继续训练、迁移学习；
最佳实践：命名规范、定期保存、验证模型；
常见问题：结构不匹配、设备不兼容、优化器状态异常。

回到用户提供的代码，在训练循环中添加模型保存逻辑（如保存最佳验证模型），并在需要时加载模型进行推理或继续训练，即可轻松实现模型的持久化管理。下次训练时，再也不用担心中断或重复劳动啦！

深度学习模型保存与装载：从训练到部署的全流程指南

一、为什么需要保存与装载模型？

二、PyTorch模型保存的核心：state_dict与完整模型

2.1 保存state_dict（推荐方式）

2.2 保存完整模型（含结构）

2.3 保存训练状态（含优化器、epoch等）

三、模型装载：从文件到推理/训练

3.1 加载state_dict（最常用）

跨设备加载（GPU→CPU）

加载部分参数（迁移学习）

3.2 加载完整模型

3.3 恢复训练（断点续训）

四、模型保存的最佳实践

4.1 命名规范

4.2 定期保存

4.3 验证保存的模型

五、常见问题与解决方案

问题1：加载模型时报错“Missing key(s) in state_dict”

问题2：加载GPU模型到CPU时报错“CUDA error: an illegal memory access was encountered”

问题3：优化器状态加载后训练效果异常

六、总结

一、为什么需要保存与装载模型？

二、PyTorch模型保存的核心：state_dict与完整模型

2.1 保存state_dict（推荐方式）

2.2 保存完整模型（含结构）

2.3 保存训练状态（含优化器、epoch等）

三、模型装载：从文件到推理/训练

3.1 加载state_dict（最常用）

跨设备加载（GPU→CPU）

加载部分参数（迁移学习）

3.2 加载完整模型

3.3 恢复训练（断点续训）

四、模型保存的最佳实践

4.1 命名规范

4.2 定期保存

4.3 验证保存的模型

五、常见问题与解决方案

问题1：加载模型时报错“Missing key(s) in state_dict”

问题2：加载GPU模型到CPU时报错“CUDA error: an illegal memory access was encountered”

问题3：优化器状态加载后训练效果异常

六、总结

二、PyTorch模型保存的核心：`state_dict`与完整模型

2.1 保存`state_dict`（推荐方式）

3.1 加载`state_dict`（最常用）

二、PyTorch模型保存的核心：`state_dict`与完整模型

2.1 保存`state_dict`（推荐方式）

3.1 加载`state_dict`（最常用）