从零开始：使用PyTorch实现手写数字识别

我们之前实现过Tensorflow实现手写数字识别，今天使用PyTorch实现手写数字识别，看看有什么不同呢。

当计算机学会"看"数字

在人工智能的世界里，计算机视觉是最具魅力的领域之一。通过PyTorch这一强大的深度学习框架，我们将在经典的MNIST数据集上，见证一个神经网络从零开始学会识别数字的全过程。本文将以通俗易懂的方式，带你走进这个看似神秘实则充满逻辑的美妙世界。

搭建开发环境

在开始训练之前，我们需要准备好三个基础要素：腾讯云HAI，腾讯云HAI，腾讯云HAI。导入必要的工具库：

import torch  # 深度学习框架核心
import torch.nn as nn  # 神经网络模块
from torchvision import datasets, transforms  # 数据处理利器

MNIST数据集解析

1. 认识手写数字数据库

MNIST数据集包含6万张训练图片和1万张测试图片，每张都是28x28像素的灰度图。这些数字由美国高中生和人口普查局员工书写，构成了计算机视觉领域的"Hello World"。

2. 数据预处理的艺术

原始图片需要经过精心处理才能被模型理解：

transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为数值矩阵
    transforms.Normalize((0.1307,), (0.3081,))  # 标准化处理
])

3. 可视化的重要性

通过Matplotlib展示样本图片，我们能直观感受数据的特征：

plt.imshow(images[0].squeeze(), cmap='gray')
plt.title(f'Label: {labels[0]}')

神经网络设计

1. 网络结构蓝图

我们设计一个全连接网络（FCN），其结构如同人类神经系统的简化版：

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()  # 将图片展开为向量
        self.fc1 = nn.Linear(28*28, 128)  # 第一隐藏层
        self.fc2 = nn.Linear(128, 64)    # 第二隐藏层
        self.fc3 = nn.Linear(64, 10)     # 输出层
        self.dropout = nn.Dropout(0.5)   # 正则化装置

• 神经元数量的选择需要平衡学习能力与过拟合风险
• Dropout层像随机关闭部分神经元，防止模型"死记硬背"

2. 信息传递机制

前向传播模拟人脑的信息处理过程：ReLU激活函数如同神经元的开关，决定是否传递信号。

def forward(self, x):
    x = self.flatten(x)  # 展平操作：将图片变为784维向量
    x = torch.relu(self.fc1(x))  # 通过第一个全连接层
    x = self.dropout(x)         # 随机屏蔽部分神经元
    x = torch.relu(self.fc2(x))  # 第二个全连接层
    return self.fc3(x)          # 最终输出10个数字的概率

让模型学会思考

1. 配置学习参数

• 损失函数：交叉熵损失（CrossEntropyLoss），衡量预测与真实的差距
• 优化器：Adam优化器，智能调节学习步伐的导航员

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

2. 训练循环解析

每个epoch都是一次完整的学习轮回：

def train(epoch):
    model.train()  # 切换至训练模式
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()  # 清空之前的梯度
        output = model(data)    # 前向传播
        loss = criterion(output, target)  # 计算损失值
        loss.backward()         # 反向传播求梯度
        optimizer.step()        # 更新网络参数

• 梯度清零避免不同批次数据的干扰
• 反向传播就像纠错老师，沿着计算链修正参数

完整代码示例

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision

from torchvision import transforms
import matplotlib.pyplot as plt

# 2. 数据准备
# 定义数据预处理：转换为Tensor并标准化（MNIST的均值和标准差）
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载训练集和测试集
train_dataset = torchvision.datasets.MNIST(
    root='./data', 
    train=True,
    download=True,
    transform=transform
)

test_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    transform=transform
)

# 创建数据加载器
train_loader = torch.utils.data.DataLoader(
    train_dataset,

    batch_size=64,
    shuffle=True
)

test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=1000,
    shuffle=False
)

# 查看数据集信息
print(f'Train samples: {len(train_dataset)}')
print(f'Test samples: {len(test_dataset)}')

# 可视化样本
images, labels = next(iter(train_loader))
plt.imshow(images[0].squeeze(), cmap='gray')
plt.title(f'Label: {labels[0]}')
plt.show()

# 3. 定义神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)
        self.dropout = nn.Dropout(0.5)
        
    def forward(self, x):
        x = self.flatten(x)
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = Net()
print(model)

# 4. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=0.001)


# 5. 训练模型
def train(epoch):

    model.train()
    running_loss = 0.0
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        

        running_loss += loss.item()
        if batch_idx % 100 == 99:

            print(f'Epoch: {epoch+1}, Batch: {batch_idx+1}, Loss: {running_loss/100:.3f}')
            running_loss = 0.0

# 6. 测试模型

def test():
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            outputs = model(data)
            _, predicted = torch.max(outputs.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    
    accuracy = 100 * correct / total
    print(f'Test Accuracy: {accuracy:.2f}%')
    return accuracy

# 7. 执行训练和测试

epochs = 5
for epoch in range(epochs):
    train(epoch)
    test()

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

针对1-9数字的测试

# 扩展测试函数，增加按数字统计的功能

def detailed_test():
    model.eval()
    class_correct = [0] * 10  # 存储每个数字的正确计数
    class_total = [0] * 10    # 存储每个数字的总样本数
    
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            
            # 遍历每个预测结果
            for label, prediction in zip(labels, predicted):
                class_total[label] += 1
                if label == prediction:
                    class_correct[label] += 1


    # 打印每个数字的准确率
    print("{:^10} | {:^10} | {:^10}".format("数字", "正确数", "准确率"))
    print("-"*33)

    for i in range(10):
        acc = 100 * class_correct[i] / class_total[i]
        print("{:^10} | {:^10} | {:^10.2f}%".format(i, class_correct[i], acc))
    
    # 可视化错误案例
    wrong_examples = []

    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        mask = predicted != labels
        wrong_examples.extend(zip(images[mask], labels[mask], predicted[mask]))
    
    # 随机展示3个错误样本

    fig, axes = plt.subplots(1, 3, figsize=(12,4))

    for ax, (img, true, pred) in zip(axes, wrong_examples[:3]):
        ax.imshow(img.squeeze(), cmap='gray')

        ax.set_title(f'True: {true}\nPred: {pred}')
        ax.axis('off')
    plt.show()

# 执行详细测试
detailed_test()

PyTorch vs TensorFlow 深度对比

1. 核心架构差异

特性	PyTorch	TensorFlow
计算图	动态图（即时执行）	静态图（需预先定义）
调试便利性	支持标准Python调试工具	需要特殊工具（tfdbg）
API设计	更接近Python原生语法	自成体系的API风格
移动端部署	支持但生态较弱	通过TF Lite有成熟解决方案

2. 相同功能的代码对比

以定义全连接层为例：

# PyTorch版
import torch.nn as nn
layer = nn.Linear(in_features=784, out_features=128)

# TensorFlow版
from tensorflow.keras.layers import Dense
layer = Dense(units=128, input_dim=784)

3. 训练流程对比

PyTorch训练循环：

for epoch in range(epochs):
    for data, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

TensorFlow训练流程：

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
model.fit(train_dataset, epochs=epochs)  # 自动完成训练循环

4. 性能对比（MNIST示例）

指标	PyTorch（CPU）	TensorFlow（CPU）
训练时间/epoch	~45秒	~50秒
内存占用	~800MB	~1GB
测试准确率	97.8-98.2%	97.5-98.0%

工具的本质

PyTorch与TensorFlow的差异，本质上是灵活性 与规范性的不同追求。就像画家选择画笔，PyTorch提供的是自由挥洒的水彩，TensorFlow则是精准可控的钢笔。理解它们的特性差异，根据项目需求选择合适的工具，才是提升开发效率的关键。无论是哪个框架，最终目标都是将数学公式转化为智能的力量。