11 从零到上手 PyTorch 的系统学习教程 - 使用 `nn.Module`

在 PyTorch 中，nn.Module 是构建神经网络的核心类。它提供了一种简单的方式来定义和组织模型的各个组成部分，包括层、参数和前向传播的计算。

1. 什么是 nn.Module

nn.Module 是 PyTorch 的一个基类，所有的神经网络模型都可以从它继承。通过扩展此类，我们可以创建自定义的神经网络层。

1.1 nn.Module 的基本结构

一个 nn.Module 通常包含以下几个部分：

• **构造函数 (__init__)**：在这里定义子模块和参数。
• **前向传播函数 (forward)**：定义输入如何通过模型进行处理。

示例

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1 = nn.Linear(10, 5)  # 输入维度为10，输出维度为5
        self.layer2 = nn.ReLU()          # ReLU 激活函数

    def forward(self, x):
        x = self.layer1(x)               # 通过第一层
        x = self.layer2(x)               # 通过激活函数
        return x

在这个例子中，我们定义了一个名为 MyModel 的模型，它包含一个线性层和一个 ReLU 激活函数。

2. 使用 nn.Module 定义复杂模型

我们不仅可以定义简单的模型，还可以通过在 __init__ 方法中添加多个层来构建更复杂的模型。

2.1 例子：创建一个简单的全连接神经网络

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(784, 128)  # 第一层：输入784（28x28）、输出128
        self.fc2 = nn.Linear(128, 10)    # 第二层：输入128，输出10（分类数）
        self.relu = nn.ReLU()            # ReLU 激活函数

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

2.2 使用模型

我们可以创建模型实例，并将数据输入模型进行前向传播。

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

# 假设我们有一个输入为 (32, 784) 的批次，代表32个样本
input_tensor = torch.randn(32, 784)

# 进行前向传播
output = model(input_tensor)
print(output.shape)  # 输出将是 (32, 10)，代表32个样本的分类输出

3. 自定义层

有时我们需要定义自定义层来实现特定功能。我们可以继承 nn.Module，并在构造函数中定义层。

示例：自定义激活函数

class CustomActivation(nn.Module):
    def forward(self, x):
        return torch.maximum(x, torch.tensor(0.0))  # 仅返回非负值

将其集成到模型中时就像使用其他层一样：

class ModelWithCustomActivation(nn.Module):
    def __init__(self):
        super(ModelWithCustomActivation, self).__init__()
        self.fc = nn.Linear(10, 5)
        self.custom_activation = CustomActivation()

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

4. 模型的训练与评估

在训练阶段，我们通常会结合优化器和损失函数来更新模型的参数。

4.1 示例：训练模型

# 假设我们使用的是交叉熵损失与随机梯度下降优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 假设我们有训练数据 train_loader
for inputs, labels in train_loader:
    optimizer.zero_grad()           # 清零之前的梯度
    outputs = model(inputs)        # 前向传播
    loss = criterion(outputs, labels)  # 计算损失
    loss.backward()                # 反向传播
    optimizer.step()               # 更新参数

4.2 评估模型性能

在评估模型时，我们通常需要计算准确度或其他指标。

# 评估模型
model.eval()  # 切换到评估模式
correct = 0
total = 0

with torch.no_grad():  # 关闭梯度计算以减少内存消耗
    for inputs, labels in test_loader:
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)  # 获取最大概率的索引
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = correct / total
print(f'Accuracy: {accuracy * 100:.2f}%')

结论

使用 nn.Module 可以轻松地构建和管理复杂的神经网络模型。通过定义 __init__ 和 forward 方法，我们能够创建可复用的组件，并便于进行训练和评估。随着对 PyTorch 的深入了解，你将能构建更加复杂的模型，进一步提升你的深度学习能力。