PyTorch实战——从零构建MNIST手写数字识别模型

1. 环境准备与数据加载

第一次接触PyTorch时，我最头疼的就是环境配置。后来发现用Anaconda管理Python环境特别省心，这里分享我的标准操作流程。先创建一个专属环境：

bash复制conda create -n pytorch_env python=3.8
conda activate pytorch_env

接着安装PyTorch全家桶，注意要根据CUDA版本选择对应命令。没有NVIDIA显卡的同学直接用CPU版本：

bash复制pip install torch torchvision tqdm matplotlib

MNIST数据集就像深度学习界的"Hello World"，包含6万张28x28像素的手写数字图片。PyTorch内置的torchvision.datasets能自动下载和处理数据，但国内用户可能会遇到下载慢的问题。这里有个小技巧——先手动下载MNIST的四个压缩文件放到./data/MNIST/raw目录下，代码就会跳过下载步骤。我用这个方式帮同事节省了半小时等待时间。

数据预处理环节有个容易踩的坑：Normalize的mean和std参数设置。MNIST是单通道图像，所以用[0.5]而不是RGB常用的[0.5, 0.5, 0.5]。transform管道这样定义：

python复制transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize(mean=[0.5], std=[0.5])
])

2. 网络结构设计实战

刚开始学CNN时，我总搞不清卷积核尺寸和特征图大小的关系。后来总结出个傻瓜公式：输出尺寸 = (输入尺寸 - kernel_size + 2*padding)/stride + 1。以MNIST为例，28x28的图片经过kernel_size=3, padding=1的卷积后，尺寸保持不变。

网络结构设计就像搭积木，这里分享我的三层CNN配方：

python复制class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.model = torch.nn.Sequential(
            torch.nn.Conv2d(1, 16, 3, 1, 1),  # 28x28 -> 28x28
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2, 2),         # 28x28 -> 14x14
            
            torch.nn.Conv2d(16, 32, 3, 1, 1), # 14x14 -> 14x14
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2, 2),         # 14x14 -> 7x7
            
            torch.nn.Conv2d(32, 64, 3, 1, 1), # 7x7 -> 7x7
            torch.nn.ReLU(),
            torch.nn.Flatten(),
            torch.nn.Linear(7*7*64, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 10)
        )

注意最后不要加Softmax！因为CrossEntropyLoss自带Softmax，重复添加会导致数值不稳定。这个坑我当年踩过，模型准确率死活上不去。

3. 训练过程优化技巧

训练循环看似简单，但藏着不少玄机。建议用tqdm包装DataLoader，实时显示进度和指标：

python复制from tqdm import tqdm

for epoch in range(EPOCHS):
    process_bar = tqdm(train_loader, unit='step')
    for step, (images, labels) in enumerate(process_bar):
        # 训练代码...
        process_bar.set_description(f"[{epoch}/{EPOCHS}] Loss:{loss.item():.4f}")

验证阶段一定要加torch.no_grad()，否则会浪费内存计算梯度。我习惯在每个epoch结束时验证一次：

python复制with torch.no_grad():
    correct = 0
    for test_imgs, test_labels in test_loader:
        outputs = net(test_imgs)
        predictions = torch.argmax(outputs, dim=1)
        correct += (predictions == test_labels).sum()
    accuracy = correct / len(test_data)

学习率设置很关键，Adam优化器默认的0.001对MNIST偏大。实测0.0005效果更好：

python复制optimizer = torch.optim.Adam(net.parameters(), lr=0.0005)

4. 模型评估与可视化

训练完成后，我习惯保存两份模型：完整模型和纯参数。前者方便直接调用，后者兼容性更好：

python复制# 保存完整模型（含结构）
torch.save(net, 'full_model.pth')
# 只保存参数
torch.save(net.state_dict(), 'params_only.pth')

可视化是检验模型的好方法。用matplotlib绘制损失和准确率曲线：

python复制plt.figure(figsize=(12,4))
plt.subplot(1,2,1)
plt.plot(train_loss_history, label='Train Loss')
plt.plot(val_loss_history, label='Validation Loss')
plt.legend()

plt.subplot(1,2,2)
plt.plot(train_acc_history, label='Train Accuracy')
plt.plot(val_acc_history, label='Validation Accuracy')
plt.legend()

想更直观可以显示预测错误的样本：

python复制wrong_samples = test_imgs[predictions != test_labels][:10]
for img in wrong_samples:
    plt.imshow(img.squeeze(), cmap='gray')
    plt.show()

5. 常见问题排查

遇到准确率卡在10%左右（随机猜测水平），通常是数据流出了问题。检查点：

1. 数据是否正常归一化（像素值应在[-1,1]区间）
2. 损失函数是否正确（分类任务用CrossEntropyLoss）
3. 优化器是否绑定正确参数

GPU内存不足时调小batch_size，但要注意batch_size影响梯度稳定性。建议从256开始尝试，最低不要小于32。

过拟合时可以考虑：

• 增加Dropout层
• 使用L2正则化
• 早停策略

python复制# 在全连接层间添加Dropout
torch.nn.Sequential(
    ...,
    torch.nn.Linear(128, 64),
    torch.nn.Dropout(0.5),
    torch.nn.ReLU(),
    ...
)

# 优化器加入L2正则化
optimizer = torch.optim.Adam(net.parameters(), weight_decay=1e-4)

6. 模型部署与应用

训练好的模型可以集成到Web应用中。用Flask搭建简易API服务：

python复制from flask import Flask, request
import torch
from PIL import Image
import io

app = Flask(__name__)
model = torch.load('model.pth', map_location='cpu')

@app.route('/predict', methods=['POST'])
def predict():
    file = request.files['image']
    img = Image.open(io.BytesIO(file.read()))
    # 预处理代码...
    with torch.no_grad():
        output = model(img)
    return str(torch.argmax(output).item())

在移动端使用时，可以考虑将模型转换为ONNX格式：

python复制dummy_input = torch.randn(1, 1, 28, 28)
torch.onnx.export(model, dummy_input, "mnist.onnx")

实际项目中，我遇到过图片尺寸不匹配的问题。解决方法是在预处理时强制resize：

python复制transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize(28),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize([0.5], [0.5])
])