PyTorch 第一个神经网络

作者likuolei

构建你的第一个 PyTorch 神经网络

本教程将引导你使用 PyTorch 构建并训练一个简单的神经网络，针对一个分类任务。我们将从头开始，创建数据、定义模型、训练并评估结果。代码简洁，适合初学者快速上手。假设你已安装 PyTorch（参考 PyTorch 官网）。

1. 任务概述

我们将构建一个全连接神经网络，处理一个模拟的二分类任务：

  • 输入：10 维随机特征（模拟数据）。
  • 输出：二分类（0 或 1）。
  • 目标：通过训练，预测样本的类别。

2. 准备环境

确保导入以下模块：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

3. 创建模拟数据集

我们生成 100 个样本，每个样本有 10 维特征，标签为 0 或 1。

# 生成模拟数据
X = torch.randn(100, 10)  # 100 个样本，10 维特征
y = torch.randint(0, 2, (100,))  # 随机二分类标签（0 或 1）

# 创建数据集和 DataLoader
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
  • 说明：
  • TensorDataset：将特征和标签配对。
  • DataLoader：支持批量加载和数据打乱。

4. 定义神经网络

我们定义一个简单的全连接神经网络，包含一个隐藏层。

class FirstNN(nn.Module):
    def __init__(self):
        super(FirstNN, self).__init__()
        self.fc1 = nn.Linear(10, 5)  # 输入 10 维，隐藏层 5 维
        self.relu = nn.ReLU()        # 激活函数
        self.fc2 = nn.Linear(5, 2)   # 输出 2 维（二分类）

    def forward(self, x):
        x = self.relu(self.fc1(x))  # 输入 -> 隐藏层 -> ReLU
        x = self.fc2(x)             # 隐藏层 -> 输出
        return x

# 实例化模型
model = FirstNN()
  • 关键点：
  • 继承 nn.Module 定义模型。
  • nn.Linear：全连接层，计算 y = Wx + b。
  • ReLU：引入非线性激活。
  • 输出 2 维，对应二分类的 logits。

5. 设置训练组件

选择损失函数、优化器，并支持 GPU（如果可用）。

# 设备选择
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# 损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 适合分类任务
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
  • 说明：
  • CrossEntropyLoss：结合 softmax 和交叉熵，适合分类。
  • SGD：随机梯度下降，lr 是学习率，momentum 加速梯度更新。

6. 训练模型

实现训练循环，包含前向传播、损失计算、反向传播和参数更新。

# 训练循环
num_epochs = 10
for epoch in range(num_epochs):
    model.train()  # 设置为训练模式
    running_loss = 0.0
    for data, target in loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()      # 清零梯度
        output = model(data)       # 前向传播
        loss = criterion(output, target)  # 计算损失
        loss.backward()            # 反向传播
        optimizer.step()           # 更新参数
        running_loss += loss.item()
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {running_loss / len(loader):.4f}")
  • 关键步骤：
  • model.train()：启用 Dropout 和 BatchNorm（如果有）。
  • optimizer.zero_grad()：防止梯度累积。
  • loss.backward()：计算梯度。
  • optimizer.step()：更新权重。

7. 评估模型

使用测试数据评估模型性能，禁用梯度计算以节省内存。

# 生成测试数据
test_data = torch.randn(10, 10).to(device)
model.eval()  # 设置为评估模式
with torch.no_grad():
    predictions = model(test_data)
    probabilities = torch.softmax(predictions, dim=1)  # 转换为概率
    predicted_classes = torch.argmax(probabilities, dim=1)  # 预测类别
    print("Predicted classes:", predicted_classes)
  • 说明：
  • model.eval()：禁用训练特定行为。
  • torch.no_grad()：停止梯度跟踪，节省内存。
  • softmax：将 logits 转换为概率。

8. 保存模型

保存训练好的模型以便复用。

# 保存模型
torch.save(model.state_dict(), 'first_nn.pth')

# 加载模型
model = FirstNN().to(device)
model.load_state_dict(torch.load('first_nn.pth'))
model.eval()

9. 完整代码

以下是整合的完整代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 数据
X = torch.randn(100, 10)
y = torch.randint(0, 2, (100,))
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# 模型
class FirstNN(nn.Module):
    def __init__(self):
        super(FirstNN, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(5, 2)

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

# 初始化
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = FirstNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 训练
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for data, target in loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {running_loss / len(loader):.4f}")

# 评估
test_data = torch.randn(10, 10).to(device)
model.eval()
with torch.no_grad():
    predictions = model(test_data)
    probabilities = torch.softmax(predictions, dim=1)
    predicted_classes = torch.argmax(probabilities, dim=1)
    print("Predicted classes:", predicted_classes)

# 保存模型
torch.save(model.state_dict(), 'first_nn.pth')

10. 结果解释

  • 训练输出：每轮（epoch）打印平均损失，损失下降表明模型在学习。
  • 预测输出：predicted_classes 显示模型对测试数据的分类结果。
  • 注意：由于数据是随机的，损失和预测可能不稳定。实际应用中应使用真实数据集（如 MNIST）。

11. 扩展建议

  • 真实数据集：尝试用 torchvision.datasets.MNIST 替换模拟数据。
  • 可视化：记录损失并绘制曲线（需要时可提供代码）。
  • 更复杂模型：添加更多层（如 nn.Conv2d 用于图像）或 Dropout 防止过拟合。
  • 超参数调优：调整学习率（lr）、批量大小（batch_size）或 epoch 数量。

12. 资源

  • 官方教程：https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html
  • 文档：https://pytorch.org/docs/stable/nn.html
  • 社区：在 X 平台搜索 #PyTorch 查看示例和讨论。

如果你需要更详细的解释、真实数据集的代码（如 MNIST）、损失曲线可视化或其他扩展功能，请告诉我！

类似文章

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注