pytorch 保存模型

在 PyTorch 中，保存和加载模型通常使用 torch.save 和 torch.load。以下是几种常见的保存模型的方式及其用法：

1. 保存整个模型

将模型结构和参数一起保存到一个文件中。
优点：加载时无需重新定义模型结构。
缺点：保存文件与代码强绑定，迁移性较差。

保存模型

import torch

# 假设有一个定义好的模型
model = YourModel()
torch.save(model, "model.pth")

加载模型

# 加载模型
model = torch.load("model.pth")
model.eval()  # 切换到评估模式

2. 仅保存模型参数（推荐）

只保存模型的参数（权重和偏置）。
优点：灵活性高，适合生产环境。
缺点：加载时需要重新定义模型结构。

保存模型参数

torch.save(model.state_dict(), "model_weights.pth")

加载模型参数

# 重新定义模型
model = YourModel()
model.load_state_dict(torch.load("model_weights.pth"))
model.eval()  # 切换到评估模式

3. 保存训练过程的检查点

保存模型参数、优化器状态以及其他信息（如 epoch 和 loss），适合训练过程中的中断恢复。

保存检查点

checkpoint = {
    "epoch": current_epoch,
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "loss": loss_value,
}
torch.save(checkpoint, "checkpoint.pth")

加载检查点

# 定义模型和优化器
model = YourModel()
optimizer = torch.optim.Adam(model.parameters())

# 加载检查点
checkpoint = torch.load("checkpoint.pth")
model.load_state_dict(checkpoint["model_state"])
optimizer.load_state_dict(checkpoint["optimizer_state"])
epoch = checkpoint["epoch"]
loss = checkpoint["loss"]

model.eval()  # 切换到评估模式

文件扩展名的选择

• .pth 或 .pt: PyTorch 官方推荐的扩展名。
• .tar: 如果保存检查点，可用 tar 格式保存多个字典。

示例综合代码

以下是一个完整的保存与加载流程示例：

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

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

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

# 假设训练过程
for epoch in range(5):
    # 模拟训练
    loss = epoch * 0.1  # 示例损失
    print(f"Epoch {epoch}, Loss: {loss}")

    # 保存检查点
    checkpoint = {
        "epoch": epoch,
        "model_state": model.state_dict(),
        "optimizer_state": optimizer.state_dict(),
        "loss": loss,
    }
    torch.save(checkpoint, "checkpoint.pth")

# 加载检查点
checkpoint = torch.load("checkpoint.pth")
model.load_state_dict(checkpoint["model_state"])
optimizer.load_state_dict(checkpoint["optimizer_state"])
print("Loaded epoch:", checkpoint["epoch"])
print("Loaded loss:", checkpoint["loss"])

这种灵活性使得 PyTorch 保存模型非常适合各种训练和部署场景。