当前位置：首页 > news >正文

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

news 2025/8/26 8:37:11

代码实现：

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}")breakdevice = '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 xmodel = NeuralNetwork().to(device)
print(model)def train(dataloader,model,loss_fn,optimizer):model.train()batch_size_num = 1for 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 = 0correct = 0with 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_batchestest_acy = correct / size * 100print(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.int64break

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→256self.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. 测试（评估）函数

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_batchestest_acy = correct / size  * 100print(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)  # 优化器：随机梯度下降，学习率1epochs = 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% 以上）。