☰
PyTorch卷积神经网络实战:从最小LeNet-5模型到训练避坑指南
2026/9/26 3:07:14 网站建设 项目流程

简介:PyTorch与卷积神经网络的实践代码包,以MNIST手写数字识别为演示项目,完整覆盖CNN的搭建、训练和评估流程。针对希望快速掌握深度学习框架应用的初学者及需要参考图像分类实现的项目开发者,代码包含数据预处理与加载、卷积层/池化层/全连接层的结构定义、损失函数与优化器配置、多轮迭代训练以及测试集准确率验证等核心环节,每部分模块化设计,易于阅读和修改。压缩包共10个文件,总大小约19.83MB,其中4个Python脚本负责模型构建、数据读取、训练与测试等逻辑,2个.pt文件为训练得到的模型权重,可直接用于推理或对比效果,另有MNIST数据集的图像与标签文件,免去手动下载数据的麻烦。目前已有1605人学习/下载,对于想在图像识别领域快速上手PyTorch的用户来说,是一份结构清晰、可直接运行的学习资料,可根据实际任务调整网络层数、卷积核大小和超参数,进一步拓展到更多分类场景。

1. 卷积神经网络遇上 PyTorch:为什么我劝你先跑通一个最小模型

先给一个反直觉的结论:在做图像分类这件事上,你的第一版卷积神经网络根本不需要多深的架构,也不需要多贵的显卡。我用 PyTorch 跑 LeNet-5 风格的网络,在 CPU 上从零训练一个十类分类器,照样能把验证集准确率推到 60% 以上。很多人一上来就搜“pytorch 卷积神经网络”,结果被各种预训练模型、注意力机制、训练技巧淹没,最后连一个能跑的脚本都没留下。这个标题真正要解决的问题,是让你用最直接的方式理解卷积神经网络在 PyTorch 里长什么样、训练流程怎么转起来、参数怎么填。

这篇文章适合两类人:一类是刚接触 PyTorch,想把环境、模型、训练串成一条线的新手;另一类是从 TensorFlow 或其他框架转过来,想快速摸清 PyTorch 的卷积层写法、踩坑点和调参习惯的熟手。我会把环境搭建、卷积层与汇聚层(也叫池化层)的参数逻辑、训练循环、常见报错和验证技巧按顺序讲清楚。每段代码你都可以直接复制到本地跑,前提是先把 PyTorch 装好——如果你还没装,下一章的环境搭建正好接上。

2. 从零搭好 PyTorch 环境:CPU 和 GPU 两条路怎么选

2.1 用 Anaconda 创建独立环境:Python 版本别乱选

我在实际项目中吃过 Python 版本太新的亏。PyTorch 的稳定版不一定第一时间支持最新的 Python 小版本,尤其你后面还可能装 torchvision、onnx 这些配套库,版本错位会让你花一下午解决依赖冲突。所以第一步,我建议用 Anaconda 建一个独立环境,Python 版本选 3.9 或 3.10,这两个版本对应 PyTorch 官方支持最稳。

conda create -n cnn python=3.10 -y conda activate cnn

这条命令创建并激活名为cnn的环境。-y表示自动确认依赖,省去交互。激活后终端提示符会从(base)变成(cnn),后续所有操作都在这个环境里,不会污染系统其他 Python。如果你已经装过 PyTorch,但环境很乱,我一般会建议直接删掉重建,conda remove -n cnn --all,别在旧环境里反复折腾。

2.2 安装 PyTorch:官方命令选版本,镜像源兜底

PyTorch 的安装命令非常依赖你的 CUDA 版本——但这个“CUDA 版本”指的是 PyTorch 编译用的 CUDA 运行时版本,不是你显卡驱动支持的 CUDA 版本。很多人搞混这两点,导致安装后 GPU 不可用。正确的做法是去 PyTorch 官网查询命令,或者用下面的固定命令。

# CUDA 12.1 对应的 GPU 版本 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

--index-url指定的是 PyTorch 官方预编译包的下载源,cu121表示 PyTorch 内部自带 CUDA 12.1 运行库。只要你的显卡驱动足够新,这个包就能调用 GPU。如果你的机器没有 NVIDIA 显卡,或者只想先在 CPU 上跑通逻辑,用这个:

# CPU 版本,体积小,适合没有独立显卡的机器 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

这里有个容易踩的坑:国内网络访问官方源可能很慢。常见做法是用清华镜像,但清华镜像的 PyPI 源里 torch 默认是 CPU 版本,带 CUDA 的包要另配。我的建议是优先用官方源,失败了再用-i https://pypi.tuna.tsinghua.edu.cn/simple安装 CPU 版本,别把两个源混在一起,否则容易出现装了 GPU 版却找不到 CUDA 库的翻车现场。下面的表可以帮你快速判断自己的情况:

使用场景安装命令倾向说明
只有 CPU,快速验证代码pip 安装 cpu 版本体积小,安装快
有 NVIDIA 显卡,跑训练pip 安装 cu121/cu118 版本需要驱动支持,显存要够
公司内网,访问外网受限本地 wheel 或镜像源 CPU 版先跑通,再换 GPU 机器

2.3 验证安装:CPU 和 GPU 状态检查不能只看 torch.cuda.is_available()

安装完成后最重要的一步是验证,但只跑torch.cuda.is_available()是不够的。我见过True照样报错的情况,因为 PyTorch 虽然能找到 CUDA 库,但显卡驱动和 CUDA 运行时版本不匹配,真正执行时才会崩溃。所以我一般会多写几行:

import torch print("PyTorch 版本:", torch.__version__) print("CUDA 是否可用:", torch.cuda.is_available()) if torch.cuda.is_available(): print("GPU 名称:", torch.cuda.get_device_name(0)) print("显存总量:", torch.cuda.get_device_properties(0).total_memory / 1024**3, "GB") # 实际跑一个小张量乘法,确认能真正执行 x = torch.randn(2, 3).cuda() print("GPU 计算测试:", (x * 2).sum().item())

这里的逻辑是:先检查is_available(),再取设备名和显存,最后用一个小张量做实际计算。如果前两步通过但最后一步报“CUDA error: no kernel image is available”,说明你的显卡架构太老,需要安装兼容的 PyTorch 版本,比如cu118或cu113。另外,如果你用的是 WSL 环境,还要确认 WSL 里能看到显卡,可以用nvidia-smi命令检查,这条命令在 WSL 里同样有效。

环境装好之后,下一步就是写模型。很多人在这一步会直接去抄复杂的网络结构,我想说的是应该先从最小网络开始。

3. 用 nn.Conv2d 搭一个能跑的 CNN:关键参数与张量形状推演

3.1 卷积层与汇聚层:参数怎么设才不玄学

卷积神经网络的核心不是“越深越好”,而是每一层的参数能不能和你输入图片的形状对上。PyTorch 里最常用的是torch.nn.Conv2d,它接收的关键参数有四个:in_channels、out_channels、kernel_size和padding。stride默认是 1,padding默认是 0。新手最容易看不懂的是输出尺寸公式:

output_size = (input_size - kernel_size + 2 * padding) / stride + 1

比如输入是 32x32 的单通道图片,kernel_size=3, padding=1, stride=1,输出仍然是 32x32,这就是“保持尺寸不变”的设置。如果你不填padding=1,输出就会缩成 30x30,后面连接全连接层时就要重新计算展平后的维度,很麻烦。代码里我会先用一个简单的打印确认形状。

还有一个你必须理解的层是汇聚层,PyTorch 里写作nn.MaxPool2d。很多资料翻译成池化层,其实“汇聚层”这个说法更贴近它的作用:把一小块区域的值汇总成一个代表值。nn.MaxPool2d(kernel_size=2, stride=2)会把尺寸减半,同时在空间上做降采样,保留最显著的特征。汇聚层没有可学习参数,但它能让网络对小幅位移更鲁棒,也能减少后续全连接层的计算量。

3.2 一个 LeNet-5 风格的最小网络

LeNet-5 是经典的手写数字识别网络,结构简单却五脏俱全,非常适合当作 PyTorch 的入门模板。我把它改成适配 32x32 三通道输入的版本,输出 10 类。下面是完整定义:

import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes=10): super().__init__() # 第一层卷积: 3通道 -> 6通道,图片尺寸保持 32x32 self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=5, padding=2) self.relu1 = nn.ReLU() # 汇聚层: 2x2窗口,步长2,输出变成 16x16 self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) # 第二层卷积: 6通道 -> 16通道,图片尺寸保持 16x16 self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, padding=2) self.relu2 = nn.ReLU() self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) # 输出变成 8x8 # 展平后送入全连接层: 16*8*8 = 1024 self.flatten = nn.Flatten() self.fc1 = nn.Linear(16 * 8 * 8, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, num_classes) def forward(self, x): x = self.relu1(self.conv1(x)) x = self.pool1(x) x = self.relu2(self.conv2(x)) x = self.pool2(x) x = self.flatten(x) x = self.fc3(self.fc2(self.fc1(x))) return x

这段代码里最关键的注释是16 * 8 * 8。16是第二层卷积的输出通道数,8x8是经过两次汇聚后的空间尺寸。如果你改了输入尺寸,比如用 64x64 的图片,这个 8 就要跟着变。我的习惯是先把模型定义好,然后专门写一行小代码打印每层输出的形状,确认无误再开始训练。另外,Flatten()默认从第二个维度开始展平,也就是保留 batch 维度,把(batch, 16, 8, 8)变成(batch, 1024),这样全连接层才能接得上。

3.3 前向传播之前先算清张量形状

下面的代码块是用来做形状验证的。我每次搭新网络都会先在随机张量上跑一次前向传播,这是最快的排错手段:

import torch model = SimpleCNN() fake_input = torch.randn(2, 3, 32, 32) # 2张32x32的3通道图片 output = model(fake_input) print("输入形状:", fake_input.shape) print("输出形状:", output.shape) # 也可以用 hook 打印中间层,但简单场景直接跑一遍即可

这段代码如果顺利执行,输出应该是torch.Size([2, 10])。2是 batch size,10是类别数。如果报维度错误,错误信息通常会告诉你期望 4 维输入但得到 3 维,或者全连接层期望输入是[-1, 1024],这时你就能定位到是某层计算错了。下面的表是输入 32x32 时各层输出形状:

层操作输出形状
输入-(2, 3, 32, 32)
conv1 + relu3->6, padding=2(2, 6, 32, 32)
pool12x2, stride=2(2, 6, 16, 16)
conv2 + relu6->16, padding=2(2, 16, 16, 16)
pool22x2, stride=2(2, 16, 8, 8)
flatten + fc1024->120->84->10(2, 10)

模型定义好了,接下来要把它训练起来。很多人卡在训练循环,不是因为代码难,而是不理解几个关键函数是干什么用的。

4. 训练循环与超参数:损失函数、优化器和学习率怎么配合

4.1 数据加载与归一化:DataLoader 的 pin_memory 和 num_workers

训练卷积神经网络,数据加载这一步决定了你的 GPU 能跑多快。我用的是 CIFAR-10 数据集,10 个类别,图片是 32x32 的三通道彩色图,正好匹配上一章的网络输入。PyTorch 的torchvision.datasets.CIFAR10可以直接下载,但要注意下载源可能在国外,网络不好时容易卡住,解决办法是手动把数据集文件放进data/目录,或者用镜像源下载后解压到对应位置。

import torchvision.transforms as T from torchvision.datasets import CIFAR10 from torch.utils.data import DataLoader transform = T.Compose([ T.ToTensor(), T.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) ]) train_dataset = CIFAR10(root='./data', train=True, transform=transform, download=True) val_dataset = CIFAR10(root='./data', train=False, transform=transform, download=True) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=4, pin_memory=True)

Normalize的参数是 CIFAR-10 数据集的全局均值和标准差,这三个数值是官方统计出来的,不要随便换成 0.5,否则模型收敛会很慢。num_workers=4表示用 4 个子进程加载数据,能把 CPU 的预处理和 GPU 计算重叠起来。pin_memory=True在 GPU 环境下能加快主机到显存的数据拷贝,但在纯 CPU 环境下没有意义,还会占用更多内存。如果你在 Windows 上跑,num_workers大于 0 可能会导致报错,这时候把它设成 0 就好,这是 Windows 上 PyTorch 的一个已知问题。

4.2 训练循环的写法:model.train() 和 optimizer.zero_grad() 不要漏

训练循环的骨架其实很固定,但我见过很多新手把optimizer.zero_grad()忘掉,导致 loss 忽高忽低,因为梯度在累积。另外,model.train()和model.eval()是两套模式,原因是部分层(比如 Dropout 和 BatchNorm)在训练和推理时的行为不同,你的网络虽然没有这些层,但养成习惯很重要。

import torch.nn.functional as F import torch.optim as optim model = SimpleCNN() optimizer = optim.Adam(model.parameters(), lr=1e-3) for epoch in range(10): model.train() running_loss = 0.0 for images, labels in train_loader: # 把梯度清零,否则上一批的梯度会累积 optimizer.zero_grad() outputs = model(images) loss = F.cross_entropy(outputs, labels) # 反向传播计算梯度 loss.backward() # 用梯度更新参数 optimizer.step() running_loss += loss.item() * images.size(0) print(f"Epoch {epoch+1}, Loss: {running_loss / len(train_dataset):.4f}")

这段代码里,F.cross_entropy其实做了两件事:对输出做 softmax,然后计算交叉熵损失。所以你的模型最后一层不需要手动加 softmax,直接输出原始 logits 就行。loss.item()拿到的是当前小批量的平均 loss,乘上images.size(0)是为了算回总量,最后除以数据集总大小得到整个 epoch 的平均损失。如果你发现 loss 在训练初期就出现nan,多半是学习率太大,或者数据归一化没做好,后面避坑章还会专门讲。

每训练一个 epoch,你应该在验证集上评估一次,防止模型过拟合训练集而不自知。验证代码要放在model.eval()模式下,并且要用torch.no_grad()关闭梯度计算,否则验证过程也会占用显存和计算资源。下面是评估代码:

def evaluate(model, val_loader): model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in val_loader: outputs = model(images) _, predicted = torch.max(outputs, dim=1) total += labels.size(0) correct += (predicted == labels).sum().item() return correct / total

torch.max(outputs, dim=1)返回两个张量,第一个是最大值,第二个是最大值所在的索引,我们只需要索引predicted。这个评估函数返回验证集准确率,可以作为调参的依据。

4.3 学习率与批次大小:先从 1e-3 和 64 开始

网络能跑之后,超参数怎么设是最玄学的部分。但你可以遵循一条可靠路线:先把 batch size 固定为 64,学习率固定为 1e-3,用 Adam 优化器跑 10 个 epoch,观察验证集准确率能否稳定上升。如果 loss 下降很慢,可以调大学习率到 3e-3;如果 loss 震荡,说明学习率太大,回调到 3e-4。下面是我常用的调试顺序:

现象调整方向说明
loss 长时间不降学习率增大到 3e-3目标太小,学习率低于 1e-4 时收敛很慢
loss 下降后回弹学习率减小到 3e-4模型在最优解附近震荡,需要更小的步长
训练集准确率高,验证集低增加 Dropout 或减小模型规模这是过拟合信号,不是学习率问题
验证集准确率上不去先检查数据归一化,再调学习率很多时候是预处理错误,不是模型问题

调参的过程中,每修改一个参数就只改一个变量,别同时动 batch size 和学习率,否则你无法判断是谁导致的变化。经验不多的人会一次性把学习率、batch、网络层数全换掉,最后模型从 50% 掉到 30%,完全找不到原因。这种“东边打一枪西边放一炮”的做法是最容易翻车的,我强调很多次也还是有人犯。

训练跑通之后,真正的拦路虎其实是各种报错。下面的避坑章,我把自己掉进去过的 5 个坑整理了一遍,现象、原因、解决方案都在里面,你可以直接对照排错。

5. 避坑:PyTorch CNN 最常见的 5 个坑(现象→原因→解决)

5.1 torch.cuda.is_available() 返回 False

现象:代码检查 CUDA 时返回 False,但nvidia-smi明明能显示显卡。

原因:最常见的情况是安装的 PyTorch 是 CPU 版本,因为很多镜像源默认把 torch 换成 CPU build。也可能显卡驱动过老,PyTorch 内置的 CUDA 运行时要求驱动版本高于某个值。还有一种情况是你在 Conda 环境里装了 CPU 版,又在系统 Python 里装了 GPU 版,调用时张冠李戴。

解决:先确认安装来源,用pip list | grep torch看 torch 版本号后面有没有+cu。没有+cu就卸载重装,命令是pip uninstall torch torchvision torchaudio,然后再用--index-url https://download.pytorch.org/whl/cu121装 GPU 版本。如果版本号确实带+cu,那多半是驱动问题,去 NVIDIA 官网更新驱动,注意笔记本用户别装错笔记本专用驱动。

5.2 前向传播报错:Expected 4D input to Conv2d

现象:模型训练时,第一层卷积就报错,信息是“Expected 4D input to Conv2d, but got 3D input”。

原因:这个错误几乎都是因为你把单张图片直接传给了模型。卷积层要求输入是(batch, channel, height, width)四维形状,即使只有一张图,也要手动增加 batch 维度。

解决:在数据加载或前向传播前检查形状,然后用images = images.unsqueeze(0)补一个维度。更规范的做法是确保 DataLoader 的batch_size大于等于 1,因为你实例化DataLoader时设了batch_size=64,按理不会出现这种问题,但在手动写测试代码时,比如从数据集中取单张图dataset[0],得到的就是(3, 32, 32),这时候最容易踩这个坑。我的习惯是写model(images.unsqueeze(0))避免性能开销和维度问题。

5.3 loss 变成 nan

现象:训练没几个迭代,loss 变成nan,准确率直接崩到 10% 左右(随机猜测水平)。

原因:最常见的原因是学习率过大,导致梯度爆炸,参数更新越过了合理范围。还有一种情况是输入数据没有归一化,比如直接把 0 到 255 的像素值喂进去,数值范围太大,配合一个不太小的学习率,很容易在反向传播时溢出。

解决:先把学习率降到 1e-4 或 3e-4,再看 loss 是否恢复正常。同时检查数据预处理,确认T.ToTensor()把像素值缩放到 0 到 1 之间,然后T.Normalize把它变成近似标准正态分布。如果仍然为 nan,可以在optimizer.step()之前加一句torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)做梯度裁剪,防止梯度范数过大,但这个只是应急手段,根本解法还是调小学习率。

5.4 训练时显存持续增长

现象:程序跑着跑着占用越来越大,最终报CUDA out of memory,但代码里每个 batch 都做了zero_grad()。

原因:一种是在循环内的 PyTorch 变量没有及时释放,尤其你用loss.item()却没有保留loss之外的计算图中间变量。更常见的是你每个 epoch 都向某个列表里追加一个loss张量,而不是追加loss.item()的浮点数。列表里保存张量会保留整个计算图,导致显存无法释放。

解决:把所有需要保存的量用.item()转成 Python 浮点数再存。训练循环里,loss.backward()后如果不再需要 loss,可以用del loss释放引用,但一般没必要。还有一种做法是在你的训练函数开头加torch.cuda.empty_cache(),但这只清理缓存,不解决真正的问题。核心原则是:不要在循环里保存任何带梯度的张量。

5.5 WSL 和国产 GPU 环境下 CUDA 版本踩坑

现象:在 WSL 里装好 PyTorch 后,torch.cuda.is_available()返回 True,但跑训练时报错 “CUDA error: no kernel image is available”。在国产 GPU 平台上更常见,明明驱动装了,PyTorch 就是识别不到。

原因:WSL 和纯 Linux 环境里的 CUDA 驱动方式不同,PyTorch 编译时针对特定 CUDA 版本生成 kernel image,如果你的显卡架构太新或太老,超出该版本的支持范围就会出这个错。国产 GPU 的情况更特殊,PyTorch 官方包没有针对它们的 kernel 适配,需要装厂商自行编译的 PyTorch 分支。

解决:WSL 用户优先安装cu118或cu121版本,如果还报错,检查显卡架构代号,搜索 PyTorch 对它的支持。国产 GPU 平台建议直接使用厂商提供的容器镜像或预编译 wheel,不要再从官网装。我见过有人在麒麟系统上硬装官方版,折腾两天无果,后来用厂商适配包半小时跑通。遇到这种平台差异,别死磕,先查官方支持列表,再用对应包。

6. 模型调优与验证:从 train_loss 到混淆矩阵,再到 ONNX 导出

最后一个阶段要解决的问题,不是让训练跑通,而是让模型真正可用。很多人在训练 loss 下降之后就停止思考,实际上还要做三件事:看验证输出、画混淆矩阵、导出模型。这里我只说一个我常用的验证技巧:在验证集上取一小批样本,把模型预测标签和真实标签打印出来,人眼比准确率数字更能发现问题。

model.eval() with torch.no_grad(): images, labels = next(iter(val_loader)) outputs = model(images) _, predicted = torch.max(outputs, 1) class_names = train_dataset.classes for i in range(8): print(f"预测: {class_names[predicted[i]]:10s} 真实: {class_names[labels[i]]:10s}")

如果你发现某个固定类别的图片总是被预测成另一个特定类别,这说明模型可能对该类别有偏差,也许是数据不平衡,也许是特征混淆。这时候再去调整训练策略,比如给损失函数加类别权重,而不是盲目改网络深度。混淆矩阵可以用 sklearn 的confusion_matrix快速生成,然后可视化,能看到哪些类别互相干扰。

模型验证满意后,下一步是导出,方便部署上线。PYTorCH 转 ONNX 是常见做法,因为 ONNX 可以跨框架运行,也能被很多推理引擎加载。下面是一个稳定的导出代码:

model.eval() dummy_input = torch.randn(1, 3, 32, 32) torch.onnx.export( model, dummy_input, "simple_cnn.onnx", input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}, opset_version=12 )

dynamic_axes让导出的模型支持可变 batch size,适合上线时候一次推理多张图。opset_version=12是兼容性和新操作权衡后的选择,太低不支持某些算子,太高可能让旧环境无法运行。导出后,建议用onnxruntime加载它,输入随机张量确认输出和 PyTorch 一致。我有个习惯是导出后比较两种框架的预测结果,除非数值精度问题,差异应该在 1e-4 以内,否则说明模型里某个算子转换出了问题。

这些验证做完,一个 PyTorch 卷积神经网络的完整闭环才算真正收尾。回想我最早踩过的坑,几乎全集中在环境选择和维度计算上,而不是网络本身。希望你读完这篇,能在自己的机器上少走几次弯路,一次把流程跑通。如果后面遇到问题,回到这里的避坑章逐条对照,大概率能找到答案。希望帮到你。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询