深度学习基础代码及解析（PyTorch实现MNIST手写数字分类）

你想要学习深度学习的实战代码与核心解析，我将以经典的MNIST手写数字分类任务为例，使用PyTorch框架实现完整流程，涵盖数据加载、模型构建、训练、评估等核心环节，代码可直接运行并附带详细解析。

一、环境准备

首先确保安装必要依赖：

# 安装PyTorch（根据系统/GPU选择对应命令，此处为CPU版本通用命令）
pip install torch torchvision

二、完整可运行代码

# 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. 超参数设置
BATCH_SIZE = 64  # 批次大小
LEARNING_RATE = 0.001  # 学习率
EPOCHS = 5  # 训练轮数
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 自动选择GPU/CPU

# 3. 数据预处理与加载
# 预处理：将PIL图像转为张量，并标准化（MNIST均值0.1307，标准差0.3081为经验值）
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 下载/加载训练集和测试集
train_dataset = datasets.MNIST(
    root='./data',  # 数据保存路径
    train=True,     # 训练集
    download=True,  # 自动下载（本地无数据时）
    transform=transform  # 应用预处理
)

test_dataset = datasets.MNIST(
    root='./data',
    train=False,    # 测试集
    download=True,
    transform=transform
)

# 数据加载器（批量加载、打乱数据、多线程读取）
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# 4. 构建深度学习模型（全连接神经网络）
class MNISTNet(nn.Module):
    def __init__(self):
        super(MNISTNet, self).__init__()
        # 定义网络层：全连接层（Linear）
        self.fc1 = nn.Linear(28 * 28, 128)  # 输入层：28*28（图像展平）→ 隐藏层128个神经元
        self.fc2 = nn.Linear(128, 64)       # 隐藏层128 → 隐藏层64
        self.fc3 = nn.Linear(64, 10)        # 输出层64 → 10个类别（0-9）
        self.relu = nn.ReLU()  # 激活函数：ReLU（解决梯度消失问题）

    def forward(self, x):
        # 定义前向传播流程（必须实现forward方法）
        x = x.view(-1, 28 * 28)  # 展平图像：(batch, 1, 28, 28) → (batch, 784)
        x = self.relu(self.fc1(x))  # 第一层：全连接+ReLU激活
        x = self.relu(self.fc2(x))  # 第二层：全连接+ReLU激活
        x = self.fc3(x)             # 第三层：全连接（输出层无需激活，后续用交叉熵损失）
        return x

# 实例化模型并移至指定设备（GPU/CPU）
model = MNISTNet().to(DEVICE)

# 5. 定义损失函数与优化器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失（适用于分类任务，内置Softmax）
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)  # Adam优化器

# 6. 模型训练函数
def train(model, train_loader, criterion, optimizer, epoch):
    model.train()  # 切换为训练模式（启用Dropout/BatchNorm等训练专属层）
    running_loss = 0.0  # 累计损失
    for batch_idx, (data, target) in enumerate(train_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()
        if batch_idx % 100 == 99:  # 每100个批次打印一次
            print(f'Epoch [{epoch+1}/{EPOCHS}], Batch [{batch_idx+1}/{len(train_loader)}], Loss: {running_loss/100:.4f}')
            running_loss = 0.0

# 7. 模型评估函数
def evaluate(model, test_loader, criterion):
    model.eval()  # 切换为评估模式（关闭Dropout/BatchNorm等，固定参数）
    test_loss = 0.0
    correct = 0  # 正确预测的样本数
    with torch.no_grad():  # 关闭梯度计算（节省内存，加速推理）
        for data, target in test_loader:
            data, target = data.to(DEVICE), target.to(DEVICE)
            output = model(data)
            test_loss += criterion(output, target).item()  # 累计测试损失
            pred = output.argmax(dim=1, keepdim=True)  # 获取预测类别（概率最大的索引）
            correct += pred.eq(target.view_as(pred)).sum().item()  # 统计正确数

    # 计算平均损失和准确率
    avg_loss = test_loss / len(test_loader)
    accuracy = 100. * correct / len(test_loader.dataset)
    print(f'\nTest set: Average loss: {avg_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')
    return accuracy

# 8. 启动训练与评估
if __name__ == '__main__':
    print(f"使用设备：{DEVICE}")
    for epoch in range(EPOCHS):
        train(model, train_loader, criterion, optimizer, epoch)
        evaluate(model, test_loader, criterion)
    # 保存模型
    torch.save(model.state_dict(), 'mnist_net.pth')
    print("模型已保存为 mnist_net.pth")

三、核心模块详细解析

1. 超参数说明

2. 数据处理流程

• transforms.Compose：组合多个预处理操作，此处先将图像转为torch.Tensor（形状：[1,28,28]，值域[0,1]），再标准化为均值0.1307、标准差0.3081，标准化可加速模型收敛。
• datasets.MNIST：PyTorch内置数据集，自动下载并加载，train=True对应60000个训练样本，train=False对应10000个测试样本。
• DataLoader：数据迭代器，实现批量加载、数据打乱（shuffle=True）、多线程读取，是PyTorch数据加载的核心工具。

3. 模型构建核心

• 必须继承nn.Module类，这是PyTorch所有模型的基类。
• __init__方法：定义网络层（全连接层nn.Linear），此处为3层全连接网络，输入层需将28×28的图像展平为784维向量。
• forward方法：定义前向传播路径，是模型的核心，决定了数据如何在网络中流动，此处使用ReLU激活函数引入非线性（无激活函数则退化为线性模型）。
• model.to(DEVICE)：将模型参数移至GPU/CPU，确保数据与模型在同一设备上（否则报错）。

4. 损失函数与优化器

• nn.CrossEntropyLoss：适用于多分类任务，内置Softmax层（将输出转为概率分布），直接接收模型原始输出（logits），无需手动添加Softmax，简化代码。
• optim.Adam：自适应学习率优化器，相比传统SGD收敛更快、更稳定，是深度学习常用优化器，传入模型参数和学习率即可。

5. 训练核心步骤（四步法）

1. optimizer.zero_grad()：清空梯度缓存，PyTorch默认梯度累积，若不清空会导致梯度错误。
2. output = model(data)：前向传播，输入数据通过模型得到预测输出。
3. loss.backward()：反向传播，根据损失值计算各参数的梯度（链式法则）。
4. optimizer.step()：优化器更新模型参数，根据梯度和学习率调整权重，最小化损失。

6. 评估关键操作

• model.eval()：切换为评估模式，关闭训练模式下的特殊层（如Dropout、BatchNorm），避免评估时参数变化。
• torch.no_grad()：上下文管理器，关闭梯度计算，减少内存占用，加速评估过程（评估无需更新参数，无需梯度）。
• output.argmax(dim=1)：获取每个样本预测概率最大的类别索引（即预测的数字），用于与真实标签比较计算准确率。

四、运行结果说明

1. 运行代码后，会自动下载MNIST数据到./data目录。
2. 每轮训练会打印批次损失，每轮训练结束后打印测试集的平均损失和准确率。
3. 5轮训练后，测试集准确率通常可达97%以上，模型保存为mnist_net.pth（可通过model.load_state_dict(torch.load('mnist_net.pth'))加载模型）。
4. 若使用GPU，训练速度会显著快于CPU。

五、核心知识点总结

1. PyTorch核心流程：数据加载→模型构建→损失/优化器定义→训练→评估。
2. 训练三要素：模型（提取特征）、损失函数（衡量误差）、优化器（更新参数）。
3. 关键模式切换：model.train()（训练）与model.eval()（评估）。
4. 梯度管理：optimizer.zero_grad()（清空梯度）、loss.backward()（计算梯度）、optimizer.step()（更新参数）。
5. 设备适配：torch.device自动适配GPU/CPU，数据与模型需统一设备。