<think>我们正在使用PyTorch保存训练好的神经网络模型。根据引用[1]和引用[4]，PyTorch提供了多种保存模型的方法。下面将详细介绍几种主要方法，并给出示例代码。
### 方法1: 保存整个模型
这种方法会保存模型的架构、权重以及优化器状态等所有信息。使用`torch.save()`函数直接保存模型对象。
```python
import torch
import torch.nn as nn
# 假设我们有一个训练好的模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc1 = nn.Linear(28*28, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = torch.flatten(x, 1)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
model = MyModel()
# ... 训练过程 ...
# 保存整个模型
torch.save(model, 'model.pth')
```
加载整个模型时，直接使用`torch.load()`：
```python
model = torch.load('model.pth')
model.eval() # 设置为评估模式
```
**注意**：这种方法在加载时依赖于原始模型类的定义。如果保存后修改了模型类，加载可能会出错。
### 方法2: 仅保存模型的状态字典（推荐）
更常用的方法是只保存模型的状态字典（`state_dict`），它包含模型的所有权重参数。这种方法更灵活，且不依赖于模型类的具体结构。
```python
# 保存模型的状态字典
torch.save(model.state_dict(), 'model_weights.pth')
```
加载时，需要先创建模型实例，然后加载状态字典：
```python
# 创建相同结构的模型
loaded_model = MyModel()
# 加载保存的权重
loaded_model.load_state_dict(torch.load('model_weights.pth'))
loaded_model.eval() # 设置为评估模式
```
### 方法3: 保存检查点（包含额外信息）
在训练过程中，我们通常希望保存检查点（checkpoint），以便在中断后恢复训练。检查点除了模型权重，还可以包含优化器状态、当前epoch、损失值等。
```python
# 假设我们还有优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# 定义检查点字典
checkpoint = {
'epoch': 10, # 当前训练轮次
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': 0.05, # 当前损失值
}
# 保存检查点
torch.save(checkpoint, 'checkpoint.pth')
```
加载检查点并恢复训练：
```python
# 创建模型和优化器
model = MyModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# 加载检查点
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
# 恢复训练
model.train() # 设置为训练模式
# ... 继续训练 ...
```
### 关键注意事项
1. **模型模式**：在保存前，确保模型处于正确的模式（训练模式或评估模式）。加载后，根据用途调用`model.train()`或`model.eval()`。
2. **设备映射**：当在GPU上训练而需要在CPU上加载时，可以使用`map_location`参数：
```python
model.load_state_dict(torch.load('model_weights.pth', map_location=torch.device('cpu')))
```
3. **文件格式**：PyTorch推荐使用`.pth`或`.pt`作为模型文件扩展名，但实际可以是任意扩展名[^1]。
4. **兼容性**：不同版本的PyTorch保存的模型可能存在兼容性问题，尽量保持环境一致。
### 完整示例流程
以下是一个完整的训练并保存模型的示例（基于引用[3]的Fashion-MNIST数据集）：
```python
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
# 定义神经网络（同引用[3]）
class QYNN(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(28 * 28, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = torch.flatten(x, start_dim=1)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# 加载数据集
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_dataset = datasets.FashionMNIST('./data', download=True, train=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
# 初始化模型和优化器
model = QYNN()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()
# 训练模型
num_epochs = 5
for epoch in range(num_epochs):
for images, labels in train_loader:
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}')
# 保存模型的状态字典（推荐方法）
torch.save(model.state_dict(), 'fashion_mnist_model.pth')
# 保存整个模型（备选）
torch.save(model, 'fashion_mnist_full_model.pth')
```
### 总结
- **推荐方法**：保存状态字典（`state_dict`）是最灵活和常用的方式。
- **检查点**：在训练过程中保存检查点，以便恢复训练。
- **部署**：保存的模型可以轻松加载用于推理（预测）或迁移学习。
§§相关问题§§
1. 如何在PyTorch中加载保存的模型进行预测？
2. 训练过程中如何定期保存检查点？
3. 如何将PyTorch模型转换为ONNX格式以用于其他框架？
4. 在加载模型时遇到“Missing key(s) in state_dict”错误如何解决？</think>### 如何在PyTorch中保存训练好的神经网络模型
在PyTorch中保存训练好的模型是深度学习工作流的关键环节，主要分为三种方法，每种方法适用于不同场景。以下是详细步骤和代码示例：
---
#### 方法1: 保存整个模型（最简单）
保存模型架构、权重和优化器状态，适合快速部署：
```python
import torch
# 定义模型（示例引用[3]）
class QYNN(torch.nn.Module):
def __init__(self):
super().__init__()
self.fc1 = torch.nn.Linear(28*28, 128)
self.fc2 = torch.nn.Linear(128, 10)
def forward(self, x):
x = torch.flatten(x, 1)
x = torch.relu(self.fc1(x))
return self.fc2(x)
# 训练模型...
model = QYNN()
# 保存整个模型（含结构和参数）
torch.save(model, 'full_model.pth') # 文件扩展名通常为.pth或.pt
```
**加载方式**：
```python
loaded_model = torch.load('full_model.pth')
loaded_model.eval() # 切换为评估模式
```
**优点**：一键保存/加载
**缺点**：文件较大，加载依赖原始类定义[^1]
---
#### 方法2: 仅保存模型参数（推荐）
仅保存权重参数，最轻量且灵活的方式：
```python
# 保存状态字典（state_dict）
torch.save(model.state_dict(), 'model_weights.pth')
```
**加载方式**：
```python
# 需先实例化相同结构的模型
new_model = QYNN()
new_model.load_state_dict(torch.load('model_weights.pth'))
new_model.eval()
```
**优点**：
- 文件小（仅参数）
- 不依赖原始类代码
- 支持跨设备加载（如GPU→CPU）
```python
# 示例：GPU模型加载到CPU
new_model.load_state_dict(torch.load('model_weights.pth', map_location=torch.device('cpu')))
```
[^4]
---
#### 方法3: 保存检查点（训练恢复）
保存训练状态（含优化器、epoch等），用于中断后恢复训练：
```python
checkpoint = {
'epoch': 10, # 当前训练轮次
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': 0.023, # 当前损失值
}
torch.save(checkpoint, 'checkpoint.pth')
```
**恢复训练**：
```python
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
current_epoch = checkpoint['epoch']
```
**适用场景**：长时间训练任务、模型微调[^1]
---
### 关键注意事项
1. **模式切换**：预测前务必调用 `model.eval()` 关闭Dropout/BatchNorm
2. **设备兼容**：使用 `map_location` 参数处理GPU/CPU转换
3. **文件格式**：
- `.pth`：PyTorch标准格式
- `.pt`：替代扩展名
- `.tar`：检查点常用格式
4. **版本控制**：不同PyTorch版本的模型可能不兼容
> **数学原理补充**：模型保存本质是存储权重矩阵 $W$ 和偏置向量 $b$。对于全连接层，前向传播计算为：
> $$\mathbf{y} = \sigma(W\mathbf{x} + b)$$
> 其中 $\sigma$ 是激活函数（如ReLU）。保存的state_dict即包含所有 $W$ 和 $b$ 的值[^2]。
---
### 完整训练到保存流程
基于Fashion-MNIST示例（引用[3]）：
```python
# 1. 数据准备
transform = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize((0.5,), (0.5,))
])
train_set = torchvision.datasets.FashionMNIST('./data', download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=32)
# 2. 模型训练
model = QYNN()
optimizer = torch.optim.Adam(model.parameters())
criterion = torch.nn.CrossEntropyLoss()
for epoch in range(10):
for images, labels in train_loader:
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 3. 保存模型参数（推荐）
torch.save(model.state_dict(), 'fashion_mnist_model.pth')
```
最低0.47元/天 解锁文章
扫一扫
1363
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
