神经网络模型搭建及手写数字识别案例

代码实现：

python 复制代码

import torch
print(torch.__version__)
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
training_data = datasets.MNIST(root='data',train=True,download=True,transform=ToTensor())
test_data = datasets.MNIST(root='data',train=False,download=True,transform=ToTensor())
from matplotlib import pyplot as plt
figure = plt.figure()
for i in range(9):
    img,labels = training_data[i + 59000]
    figure.add_subplot(3,3,i + 1)
    plt.title(labels)
    plt.axis('off')
    plt.imshow(img.squeeze(),cmap="gray")
plt.show()
train_dataloader = DataLoader(training_data,batch_size=64)
test_dataloader = DataLoader(test_data,batch_size=64)

for X,y in test_dataloader:
    print(f"Shape of X[N,C,H,W]:{X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
print(f"Using {device} device")

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.hidden1 = nn.Linear(28*28,128)
        self.hidden2 = nn.Linear(128,256)
        self.out = nn.Linear(256,10)
    def forward(self,x):
        x = self.flatten(x)
        x = self.hidden1(x)
        x = torch.sigmoid(x)
        x = self.hidden2(x)
        x = torch.sigmoid(x)
        x = self.out(x)
        return x

model = NeuralNetwork().to(device)
print(model)

def train(dataloader,model,loss_fn,optimizer):
    model.train()
    batch_size_num = 1
    for X ,y in dataloader:
        X,y = X.to(device),y.to(device)
        pred = model.forward(X)
        loss = loss_fn(pred,y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        loss_value = loss.item()
        if batch_size_num % 100 ==0:
            print(f"loss:{loss_value:>7f} [number:{batch_size_num}]")
        batch_size_num +=1
def test(dataloader,model,loss_fn):
    size = len(dataloader.dataset)
    num_batches= len(dataloader)
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for X ,y in dataloader:
            X,y = X.to(device),y.to(device)

            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_pj_loss = test_loss / num_batches
    test_acy = correct / size * 100
    print(f"Avg loss: {test_pj_loss:>7f} \n Accuray: {test_acy:>5.2f}%")
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=0.01)
# train(train_dataloader,model,loss_fn,optimizer)
# test(test_dataloader,model,loss_fn)
i=10
for j in range(i):
    print(f"Epoch {j+1}\n----------")
    train(train_dataloader, model,loss_fn,optimizer)
print("Done!")
test(test_dataloader,model,loss_fn)

这段代码是通过PyTorch 实现的 MNIST 手写数字识别神经网络，包含数据加载、可视化、模型构建、训练和评估的完整流程。下面分步骤解析：

1. 库导入

复制代码

import torch  # 导入PyTorch核心库
print(torch.__version__)  # 打印PyTorch版本
from torch import nn  # 导入神经网络模块
from torch.utils.data import DataLoader  # 导入数据加载工具
from torchvision import datasets  # 导入计算机视觉数据集
from torchvision.transforms import ToTensor  # 导入图像转张量的工具
from matplotlib import pyplot as plt  # 导入可视化库

核心库：torch是 PyTorch 的主库，nn用于构建神经网络，DataLoader用于批量加载数据。

数据集：torchvision.datasets提供了 MNIST 等经典数据集，ToTensor将图像（PIL 格式）转换为 PyTorch 张量（便于计算）。

可视化：matplotlib用于展示数据样本。

2. 数据加载

复制代码

# 加载MNIST训练集（60000张图片），若本地没有则自动下载，并用ToTensor转换
training_data = datasets.MNIST(
    root='data',  # 数据存储路径
    train=True,   # 标记为训练集
    download=True,  # 自动下载
    transform=ToTensor()  # 转换为张量
)

# 加载MNIST测试集（10000张图片），参数含义同上
test_data = datasets.MNIST(
    root='data',
    train=False,  # 标记为测试集
    download=True,
    transform=ToTensor()
)

MNIST 是手写数字数据集（0-9），每张图片为 28×28 像素的灰度图，常用于图像识别入门。

train=True对应训练集（用于模型学习），train=False对应测试集（用于评估模型性能）。

3. 数据可视化

复制代码

figure = plt.figure()  # 创建画布
for i in range(9):  # 展示9张图片
    img, labels = training_data[i + 59000]  # 取训练集中第59000+1到59000+9张图片
    figure.add_subplot(3, 3, i + 1)  # 3行3列布局
    plt.title(labels)  # 显示标签（真实数字）
    plt.axis('off')  # 关闭坐标轴
    plt.imshow(img.squeeze(), cmap="gray")  # 显示灰度图（squeeze移除多余维度）
plt.show()  # 展示图片

作用：直观查看数据样本，确认数据加载正确（图片与标签是否匹配）。

img.squeeze()：原始图像张量形状为(1,28,28)（1 个通道，28×28 像素），squeeze()移除通道维度，变为(28,28)便于显示。

4. 数据批量处理

复制代码

# 将数据集转换为可迭代的批量数据加载器
train_dataloader = DataLoader(training_data, batch_size=64)  # 训练集，每批64个样本
test_dataloader = DataLoader(test_data, batch_size=64)  # 测试集，每批64个样本

# 打印一个测试批次的数据形状
for X, y in test_dataloader:
    print(f"Shape of X[N,C,H,W]: {X.shape}")  # 输出：[64,1,28,28]
    print(f"Shape of y: {y.shape} {y.dtype}")  # 输出：[64] torch.int64
    break

DataLoader的作用：将数据集拆分为多个批次（batch_size=64），支持并行加载，提高训练效率。

数据形状说明：

X（图像）：[N, C, H, W]，其中N=64（批次大小）、C=1（灰度图通道数）、H=28、W=28（图像尺寸）。

y（标签）：[64]，每个元素是 0-9 的整数（表示图像对应的数字）。

5. 设备选择

复制代码

# 优先使用GPU（cuda），其次苹果芯片GPU（mps），最后CPU
device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
print(f"Using {device} device")

作用：选择计算设备，GPU（cuda/mps）的并行计算能力可大幅加速神经网络训练，CPU 速度较慢。

6. 神经网络模型定义

复制代码

class NeuralNetwork(nn.Module):  # 继承nn.Module（PyTorch神经网络基类）
    def __init__(self):
        super().__init__()  # 初始化父类
        self.flatten = nn.Flatten()  # 展平层：将28×28图像转为1维向量
        self.hidden1 = nn.Linear(28*28, 128)  # 全连接层1：784→128（输入784=28×28）
        self.hidden2 = nn.Linear(128, 256)  # 全连接层2：128→256
        self.out = nn.Linear(256, 10)  # 输出层：256→10（10个类别，对应0-9）

    def forward(self, x):  # 定义前向传播（必须实现）
        x = self.flatten(x)  # 展平：(64,1,28,28)→(64,784)
        x = self.hidden1(x)  # 第一层计算：(64,784)→(64,128)
        x = torch.sigmoid(x)  # 激活函数：引入非线性
        x = self.hidden2(x)  # 第二层计算：(64,128)→(64,256)
        x = torch.sigmoid(x)  # 激活函数
        x = self.out(x)  # 输出层：(64,256)→(64,10)
        return x

# 实例化模型并移动到指定设备
model = NeuralNetwork().to(device)
print(model)  # 打印模型结构

模型结构：3 层全连接网络（2 个隐藏层 + 1 个输出层），通过nn.Linear定义线性变换，sigmoid激活函数引入非线性（使模型能拟合复杂关系）。

forward方法：定义数据在网络中的流动过程，是模型计算的核心。

7. 训练函数

复制代码

def train(dataloader, model, loss_fn, optimizer):
    model.train()  # 设为训练模式（启用 dropout/batchnorm等训练特有的层）
    batch_size_num = 1  # 批次计数器
    for X, y in dataloader:  # 遍历每个批次
        X, y = X.to(device), y.to(device)  # 数据移到设备（GPU/CPU）
        
        # 前向传播：计算预测值
        pred = model.forward(X)
        # 计算损失（预测值与真实标签的差距）
        loss = loss_fn(pred, y)

        # 反向传播+参数更新
        optimizer.zero_grad()  # 清空上一轮梯度
        loss.backward()  # 反向传播计算梯度
        optimizer.step()  # 优化器更新模型参数

        # 每100个批次打印一次损失
        loss_value = loss.item()  # 取出损失值（转为Python数值）
        if batch_size_num % 100 == 0:
            print(f"loss: {loss_value:>7f} [number: {batch_size_num}]")
        batch_size_num += 1

核心逻辑：通过 "前向传播计算损失→反向传播求梯度→优化器更新参数" 的循环，让模型逐步学习数据规律。

model.train()：启用训练模式（部分层如 Dropout 在训练和测试时行为不同）。

梯度清空：optimizer.zero_grad()避免梯度累积，保证每轮梯度计算独立。

8. 测试（评估）函数

python 复制代码

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches= len(dataloader)
    model.eval()  # 设为评估模式（关闭 dropout/batchnorm等训练特有的层）
    test_loss = 0  # 总损失
    correct = 0  # 正确预测数
    with torch.no_grad():  # 关闭梯度计算（节省内存，加速评估）
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)  # 前向传播（不计算梯度）
            test_loss += loss_fn(pred, y).item()  # 累加损失
            # 计算正确预测数：pred.argmax(1)取预测概率最大的类别，与y比较
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
   
    # 计算平均损失和准确率
    test_pj_loss = test_loss /  num_batches
    test_acy = correct / size  * 100
    print(f"Avg loss: {test_pj_loss:>7f} \n Accuracy: {test_acy:>5.2f}%")

作用：评估模型在测试集上的性能（泛化能力），不更新参数。

model.eval()：关闭训练特有的层（如 Dropout），确保评估稳定。

torch.no_grad()：关闭梯度计算，减少内存占用，加速评估。

准确率计算：pred.argmax(1)获取每个样本预测的最大概率类别，与真实标签y比较，统计正确比例。

9. 训练与评估执行

复制代码

loss_fn = nn.CrossEntropyLoss()  # 损失函数：适合多分类问题（内置Softmax）
optimizer = torch.optim.SGD(model.parameters(), lr=1)  # 优化器：随机梯度下降，学习率1

epochs = 10  # 训练轮数（完整遍历训练集10次）
for j in range(epochs):
    print(f"Epoch {j+1}\n----------")
    train(train_dataloader, model, loss_fn, optimizer)  # 训练一轮
print("Done!")
test(test_dataloader, model, loss_fn)  # 训练完成后在测试集评估

损失函数：CrossEntropyLoss是多分类任务的常用损失，结合了nn.LogSoftmax和nn.NLLLoss，直接接收原始输出（无需手动加 Softmax）。

优化器：SGD（随机梯度下降）用于更新模型参数，lr=1是学习率（控制参数更新幅度）。

训练轮数（epochs=10）：模型会完整遍历训练集 10 次，逐步降低损失、提高准确率。

最终评估：训练完成后，在测试集上输出平均损失和准确率（通常能达到 90% 以上）。

总结

这段代码完整实现了一个基于全连接网络的 MNIST 手写数字识别流程，涵盖数据加载、可视化、模型构建、训练和评估。核心逻辑是通过反向传播算法，让模型从数据中学习 "图像像素→数字类别" 的映射关系，最终实现对未知手写数字的识别。