Pytorch - Model

使用示例

概念

在 PyTorch 中，Model 是指继承自 torch.nn.Module 的类，表示一个神经网络模型。为了更好地理解如何定义、保存、加载和使用预训练模型，我们可以逐一详细介绍这些内容。

模型的定义

在 PyTorch 中，模型的定义通常包括以下几个步骤：

定义 __init__ 方法**

__init__ 方法用于初始化模型的各个层和模块。通常，你会在这里定义网络的结构，包括各个神经网络层（如：全连接层、卷积层、池化层等）。

定义 forward 方法**

forward 方法定义了前向传播的过程。它接受输入张量，按照定义的网络结构将数据流过每一层，最终输出结果。

模型的定义范例

下面是一个简单的多层感知机（MLP）模型定义的例子：

详情
import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        # 定义全连接层
        self.fc1 = nn.Linear(784, 128)  # 输入784维，输出128维
        self.relu = nn.ReLU()           # ReLU激活函数
        self.fc2 = nn.Linear(128, 10)   # 输入128维，输出10维（用于分类10个类别）

    def forward(self, x):
        # 定义前向传播过程
        x = self.fc1(x)   # 通过第一个全连接层
        x = self.relu(x)   # 通过ReLU激活
        x = self.fc2(x)    # 通过第二个全连接层
        return x

在上述例子中：

  • __init__ 方法中定义了两层全连接层和一个激活函数。
  • forward 方法则实现了数据从输入到输出的流动。

模型的训练

训练循环

一旦模型定义完成，我们就可以开始训练模型。训练过程通常包括以下几个步骤：

  • 定义损失函数（Loss Function）
  • 定义优化器（Optimizer）
  • 设置训练数据和目标（DataLoader）
  • 在每个训练周期中，进行正向传播、计算损失、反向传播、更新参数

下面是一个简单的训练代码示例：

详情
import torch.optim as optim

# 假设我们有输入数据 x 和目标输出 y
x = torch.randn(64, 784)  # 批量大小64，每个输入784个特征
y = torch.randint(0, 10, (64,))  # 假设目标标签为 0 到 9 之间的整数

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

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练过程
for epoch in range(10):  # 假设训练10个周期
    model.train()  # 设置模型为训练模式
    optimizer.zero_grad()  # 清空之前的梯度

    # 正向传播
    outputs = model(x)
    loss = criterion(outputs, y)  # 计算损失

    # 反向传播
    loss.backward()
    optimizer.step()  # 更新模型参数

    print(f'Epoch [{epoch+1}/10], Loss: {loss.item():.4f}')

模型的保存与加载

保存模型

在 PyTorch 中，模型的保存有两种方式：

  • 保存整个模型（包括结构和参数）
  • 仅保存模型的参数（推荐方式）

保存整个模型：

torch.save(model, 'model.pth')

保存模型的参数：

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

保存模型参数的方法较为常用，因为它更灵活。在加载模型时，你可以重新定义模型结构，然后加载模型参数。

加载模型

加载整个模型：

# 直接加载整个模型
model = torch.load('model.pth')
model.eval()  # 切换到评估模式

加载模型的参数：

# 重新定义模型
model = SimpleModel()
# 加载模型参数
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()  # 切换到评估模式

保存和加载模型的优缺点

  • 保存整个模型：

    • 优点：简单，直接保存模型结构和参数。
    • 缺点：保存的模型依赖于代码实现，不能轻松迁移或共享给其他人。

  • 保存模型参数：

    • 优点：灵活，适合共享和迁移。可以在不同环境中加载相同的模型结构。
    • 缺点：需要你在加载时重新定义模型结构。

预训练模型

预训练模型是指已经在某个大数据集上训练过的模型（例如：ImageNet、COCO 等）。这些模型通常已经学会了某些特定领域的特征，因此可以通过迁移学习加速你自己的任务。

PyTorch 提供了很多预训练的模型，通常来自于 torchvision.models 模块。常见的预训练模型包括 ResNet、VGG、Inception、MobileNet 等。

使用预训练模型

你可以轻松地加载预训练模型并进行微调或特征提取。

加载预训练的 ResNet50：

import torchvision.models as models

# 加载预训练的 ResNet50 模型
model = models.resnet50(pretrained=True)

# 如果你只想提取特征，或者微调最后一层，可以修改最后一层
model.fc = nn.Linear(model.fc.in_features, 10)  # 假设我们做10分类任务
  • pretrained=True 表示加载预训练的权重。如果你不想使用预训练权重，可以设置为 pretrained=False。

微调预训练模型

微调（Fine-Tuning）是指在已有的预训练模型基础上，针对自己的任务数据进行训练。通常，我们会冻结大部分层的权重，只训练最后几层。

# 冻结所有层的参数
for param in model.parameters():
    param.requires_grad = False

# 解冻最后一层的参数
for param in model.fc.parameters():
    param.requires_grad = True

# 选择优化器时，确保只优化未冻结的参数
optimizer = optim.SGD(model.fc.parameters(), lr=0.01)

模型评估与推断

训练完成后，我们通常会使用模型进行评估和推断。

评估模式

当进行模型评估时，应该将模型设置为评估模式（model.eval()），这会关闭 Dropout 层和 BatchNorm 层的训练行为。

model.eval()  # 设置为评估模式
with torch.no_grad():  # 在评估阶段，不需要计算梯度
    outputs = model(x_test)  # 对测试数据进行推断
    predicted_class = outputs.argmax(dim=1)  # 获取预测的类别

推断

推断通常是在评估模式下进行的。我们使用 torch.no_grad() 来关闭梯度计算，因为我们只关心模型的输出而不需要更新模型参数。

总结

  • 模型定义：通过继承 torch.nn.Module，你可以自定义神经网络结构，定义 __init__ 和 forward 方法来实现模型的层和前向传播。
  • 模型训练：使用 torch.optim 定义优化器，使用 nn.CrossEntropyLoss 等定义损失函数，然后进行训练和反向传播。
  • 保存和加载模型：保存模型时，你可以选择保存整个模型或仅保存模型的参数。加载模型时，通常只加载参数并重新定义网络结构。
  • 预训练模型：使用 PyTorch 提供的预训练模型（例如 ResNet、VGG 等），你可以直接加载并进行微调以适应你的任务。
最近更新:
Contributors: klc407073648