用Pytorch保存模型与参数（已训练好神经网络模型）

编程语言 2025-04-09 19:19:09 阅读次数: 0

1. 前言

2. 保存模型的三种主要方法

1. 前言

在深度学习的实践中，训练一个神经网络模型通常需要花费大量的时间和计算资源。因此，保存训练好的模型是一个非常重要的步骤，它不仅能让我们在后续直接加载和使用模型，还能为模型的迁移学习、部署和共享提供便利。PyTorch 作为一款流行的深度学习框架，提供了多种保存和加载模型的方法。本文将详细讲解如何在 PyTorch 中保存训练好的神经网络模型，并探讨不同方法的适用场景和注意事项。

2. 保存模型的三种主要方法

2.1 保存整个模型

保存整个模型是最直接的方法，它会将模型的结构和参数一起保存到一个文件中。这种方法适合简单的场景，尤其是当你不需要对模型进行进一步修改时。

步骤：

训练模型：首先确保你的模型已经完成训练。
使用 torch.save() 保存模型：将模型对象直接传递给 torch.save() 函数。

import torch
import torch.nn as nn

# 定义一个简单的神经网络
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(10, 2)
    
    def forward(self, x):
        return self.fc(x)

# 创建模型实例并训练（假设已经训练完成）
model = SimpleNet()

# 保存整个模型
torch.save(model, 'whole_model.pth')

加载模型：

加载保存的整个模型时，可以直接使用 torch.load() 函数。

# 加载整个模型
loaded_model = torch.load('whole_model.pth')

# 测试模型
input_data = torch.randn(1, 10)
output = loaded_model(input_data)
print(output)

2.2 只保存模型参数（推荐）

在实际应用中，通常推荐只保存模型的参数（state_dict），而不是整个模型。这是因为保存参数更加灵活，可以方便地将参数迁移到其他模型结构中。

步骤：

训练模型：确保模型已经完成训练。
提取模型参数：使用 model.state_dict() 获取模型的参数。
保存参数：将参数保存到文件中。

# 提取并保存模型参数
torch.save(model.state_dict(), 'model_params.pth')

加载参数：

加载参数时，需要先定义模型结构，然后将参数加载到模型中。

# 定义相同的模型结构
loaded_model = SimpleNet()

# 加载参数
loaded_model.load_state_dict(torch.load('model_params.pth'))

# 测试模型
input_data = torch.randn(1, 10)
output = loaded_model(input_data)
print(output)

2.3 保存训练状态（包括优化器）

在某些情况下，你可能希望保存模型的训练状态，包括优化器的状态（如学习率、动量等）。这在需要从中断处继续训练时非常有用。

步骤：

定义模型和优化器：确保模型和优化器已经初始化。
训练模型：在训练过程中，可以定期保存训练状态。
保存训练状态：将模型参数和优化器状态一起保存。

# 定义模型和优化器
model = SimpleNet()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 假设已经训练了几轮
for epoch in range(10):
    # 训练代码
    pass

# 保存训练状态
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, 'training_state.pth')

加载训练状态：

加载训练状态时，可以从中断处继续训练。

# 加载训练状态
checkpoint = torch.load('training_state.pth')

# 定义相同的模型和优化器
model = SimpleNet()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 加载状态
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch'] + 1

# 继续训练
for epoch in range(start_epoch, 20):
    # 训练代码
    pass

3. 保存和加载模型的注意事项

设备兼容性：保存的模型在加载时可能会遇到设备（CPU/GPU）不匹配的问题。可以通过 map_location 参数指定加载到的设备。
```
loaded_model = torch.load('model_params.pth', map_location=torch.device('cpu'))
```
模型结构一致性：只保存参数时，加载时必须确保模型结构与保存时一致，否则会报错。
文件格式：PyTorch 默认使用 .pth 或 .pt 作为保存文件的扩展名，但你可以根据需要选择其他扩展名。
定期保存：在长时间训练中，建议定期保存模型状态，以防止因意外中断导致的训练成果丢失。