AI深度学习入门教程:用PyTorch跑通MNIST手写数字识别
深度学习入门最容易卡住的地方不是理论,而是环境跑不起来、代码报错、训练结果看不懂。
这篇教程直接带你用 PyTorch 跑通 MNIST 手写数字识别,每一步都有命令和预期结果,适合刚学完 Python 基础、想动手理解神经网络的新手。
先确认你要装什么、在哪运行
MNIST 是一个 7 万张 28x28 灰度手写数字图片的数据集,常被当作深度学习的“Hello World”。
本文使用 PyTorch 框架和 CPU 或单卡 GPU 均可运行,建议 Python 3.9 以上版本。
如果你还没有安装 Python,先到 python.org 下载安装,并在终端确认版本:
python --version
预期输出类似 Python 3.10.12。
版本过低可能导致后续依赖安装失败。
创建独立虚拟环境,避免污染系统包
在项目目录下执行以下命令,创建并激活虚拟环境:
python -m venv dl_env
source dl_env/bin/activate # Linux / macOS
# Windows 使用: dl_env\Scripts\activate
激活后命令行前面会出现 (dl_env) 标识。
接着安装 PyTorch 和 torchvision。
CPU 版本安装命令如下:
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
如果你有 NVIDIA 显卡并已安装对应驱动,可以到 PyTorch 官网选择匹配 CUDA 版本的安装命令,不要直接复制旧教程里的固定版本号。
安装完成后验证:
python -c "import torch; print(torch.__version__)"
能正常打印版本号即表示安装成功。
加载数据并定义网络结构
新建文件 train_mnist.py,先写入数据加载部分:
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_set = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_set = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
test_loader = DataLoader(test_set, batch_size=1000, shuffle=False)
第一次运行会自动下载数据到 ./data 目录。
如果下载超时,检查网络或手动下载数据集放到对应目录。
接着定义一个简单的卷积神经网络:
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
self.pool = nn.MaxPool2d(2)
self.fc1 = nn.Linear(64 * 7 * 7, 128)
self.fc2 = nn.Linear(128, 10)
self.relu = nn.ReLU()
def forward(self, x):
x = self.pool(self.relu(self.conv1(x)))
x = self.pool(self.relu(self.conv2(x)))
x = x.view(x.size(0), -1)
x = self.relu(self.fc1(x))
return self.fc2(x)
这段代码里,卷积层负责提取图像特征,池化层缩小尺寸,全连接层输出 0 到 9 十个类别的分数。
新手先理解输入输出形状即可,不必一开始推导每个公式。
训练循环与结果验证
继续在文件末尾加入训练和测试逻辑:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(3):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
if batch_idx % 200 == 0:
print(f'Epoch {epoch} batch {batch_idx} loss {loss.item():.4f}')
model.eval()
correct = 0
total = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
pred = output.argmax(dim=1)
correct += (pred == target).sum().item()
total += target.size(0)
print(f'Epoch {epoch} test accuracy: {correct / total:.4f}')
运行 python train_mnist.py,正常情况下每个 epoch 结束后会打印测试准确率,3 个 epoch 后一般能达到 98% 以上。
准确率是模型在没见过的测试集上答对的比例,比训练损失更能说明效果。
常见报错和避坑提醒
ModuleNotFoundError: No module named 'torch':虚拟环境没激活,或安装时用了系统 Python。重新激活dl_env后再执行 pip 安装。- 下载 MNIST 卡住或超时:可以设置代理,或从镜像站手动下载四个压缩包放到
./data/MNIST/raw目录,再重新运行。 - GPU 显存不足:把
batch_size从 64 降到 32 或 16;如果仍报错,先用 CPU 跑通流程。 - 准确率一直不上升:检查是否误把
optimizer.zero_grad()删掉,或学习率设置过大。先用本文参数跑通再调。 - 不要一开始就改网络层数:入门阶段先保证数据加载、训练循环、评估三部分正确,再加复杂结构。
跑通之后怎么继续学
这个模型能跑通,说明你已经走完了深度学习最核心的流程:准备数据、定义网络、计算损失、反向传播更新参数、评估效果。
随后可以尝试保存模型权重、把测试集里的错误样本打印出来观察,或者换用 FashionMNIST 数据集对比结果。
遇到报错时优先回看虚拟环境、设备类型和 batch 大小这三处,大多数入门问题都集中在这里。