PyTorch入门:构建简单神经网络实现图像分类

下面将逐步指导如何使用PyTorch构建一个全连接神经网络,在MNIST手写数字数据集上实现图像分类。MNIST包含28×28像素的灰度手写数字图像(0-9共10类)。


1. 环境准备

首先导入所需库:

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader


2. 数据预处理

使用torchvision加载MNIST数据,并进行标准化处理:

# 数据转换:转换为Tensor并归一化到[-1,1]
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载数据集
train_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_data = datasets.MNIST(root='./data', train=False, transform=transform)

# 创建数据加载器
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64)


3. 构建神经网络模型

定义一个包含两个隐藏层的全连接网络:

class NeuralNet(nn.Module):
    def __init__(self):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(28*28, 128)  # 输入层:784像素 -> 128神经元
        self.fc2 = nn.Linear(128, 64)     # 隐藏层:128 -> 64
        self.fc3 = nn.Linear(64, 10)      # 输出层:64 -> 10类别
        
    def forward(self, x):
        x = x.view(-1, 28*28)            # 展平图像 [batch, 784]
        x = torch.relu(self.fc1(x))       # ReLU激活
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)                   # 输出层无需激活
        return x

model = NeuralNet()


4. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()         # 交叉熵损失
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam优化器


5. 训练模型
epochs = 5  # 训练轮数

for epoch in range(epochs):
    for images, labels in train_loader:
        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')


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

with torch.no_grad():  # 禁用梯度计算
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)  # 获取预测类别
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'测试准确率: {100 * correct / total:.2f}%')


关键概念解释
  1. 张量(Tensor):PyTorch基础数据结构,支持GPU加速
  2. 自动微分(Autograd):通过loss.backward()自动计算梯度
  3. 计算图:动态构建计算路径,便于梯度传播
  4. 激活函数:ReLU引入非线性,公式:$f(x) = \max(0, x)$
  5. 交叉熵损失:分类任务常用损失函数,定义为: $$ \mathcal{L} = -\sum_{i=1}^{N} y_i \log(\hat{y}_i) $$ 其中$y_i$为真实标签,$\hat{y}_i$为预测概率

改进建议
  1. 使用卷积神经网络(CNN)提升图像识别精度
  2. 添加Dropout层防止过拟合
  3. 使用学习率调度器(如StepLR
  4. 尝试其他优化器(如RMSprop

完整代码约50行,可在GPU环境下运行(将张量移至GPU:images = images.to(device))。

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐