PyTorch保存训练好的模型

原创

mob64ca12d2dee8 2024-09-29 03:59:08 ©著作权

©著作权归作者所有：来自51CTO博客作者mob64ca12d2dee8的原创作品，请联系作者获取转载授权，否则将追究法律责任

PyTorch模型保存与加载指南

在深度学习的训练过程中，保存和加载训练好的模型是一个非常重要的步骤。PyTorch提供了简便的方法来实现这一过程。本文将详细介绍如何使用PyTorch保存和加载训练好的模型。

工作流程

以下是实现保存与加载模型的基本步骤：

步骤	描述
1. 定义模型	创建一个PyTorch神经网络模型
2. 训练模型	使用训练数据对模型进行训练
3. 保存模型	将训练好的模型保存到文件
4. 加载模型	从文件加载模型并进行推理或进一步训练

每一步的详细操作

1. 定义模型

首先，我们需要定义一个简单的神经网络模型。这里我们用一个线性层的示例。

import torch
import torch.nn as nn

# 定义一个简单的神经网络模型
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(10, 2)  # 输入维度10, 输出维度2

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

# 创建模型实例
model = SimpleNet()

2. 训练模型

接下来，让我们假设我们已经有了训练数据，下面是如何训练模型的简单示例：

# 假设我们有训练数据
data = torch.randn((5, 10))  # 5个样本，每个样本10个特征
labels = torch.randint(0, 2, (5,))  # 生成随机的标签

# 使用交叉熵损失和Adam优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 模型训练的示例
for epoch in range(100):  # 训练100轮
    optimizer.zero_grad()   # 清空梯度
    outputs = model(data)   # 进行前向传播
    loss = criterion(outputs, labels)  # 计算损失
    loss.backward()         # 反向传播
    optimizer.step()        # 更新参数

3. 保存模型

训练完成后，我们可以使用以下代码将模型保存到文件中。

# 保存模型的状态字典
torch.save(model.state_dict(), 'simple_net.pth')

4. 加载模型

当我们需要使用保存的模型时，我们可以轻松加载它并进行推理。

# 创建新的模型实例
loaded_model = SimpleNet()
# 加载模型的状态字典
loaded_model.load_state_dict(torch.load('simple_net.pth'))
# 将模型设置为评估模式
loaded_model.eval（)

类图

以下是模型类的类图：

classDiagram
class SimpleNet {
    +__init__()
    +forward(x)
}

旅行图

下图展示了保存与加载模型过程中的重要步骤：

journey
    title PyTorch模型保存与加载
    section 模型训练
      定义模型          : 2: 开发者
      训练模型          : 4: 开发者
    section 模型保存与加载
      保存模型          : 5: 开发者
      加载模型          : 3: 开发者
      使用模型          : 4: 研究者