资讯动态

PyTorch实战:从零搭建卷积神经网络CNN,完成CIFAR-10图像分类

发布时间:2026/8/22 6:55:18 来源:尧图企业网站定制
1. 项目概述为什么选择PyTorch搭建你的第一个CNN如果你刚接触深度学习面对TensorFlow、PyTorch这些框架不知从何下手或者已经看过一些理论但一写代码就报错感觉理论和实践之间隔着一道鸿沟那么这篇内容就是为你准备的。我将带你从零开始用PyTorch亲手搭建一个能真正跑起来的卷积神经网络CNN模型。这不是一个简单的“Hello World”而是一个完整的、可以用于图像分类任务的实战项目涵盖了从环境搭建、数据准备、模型定义、训练调试到评估保存的全流程。为什么是PyTorch这几年它的流行度有目共睹尤其是在研究和快速原型开发领域。相比其他框架PyTorch的“动态计算图”特性让它写起来像用NumPy一样直观调试起来也方便得多。你写的每一行代码都直接对应着一次计算哪里出错一目了然。对于初学者来说这种即时反馈的学习体验至关重要。而CNN作为计算机视觉的基石从人脸识别到医学影像分析无处不在理解它是进入深度学习视觉领域的必经之路。通过这个项目你不仅能学会PyTorch的基本操作更能深刻理解CNN的各个组件卷积层、池化层、全连接层是如何协同工作的把书本上的方块图变成可以运行的代码。2. 环境准备与工具选型打造稳定的开发基石在开始写第一行模型代码之前一个稳定、兼容的环境是成功的一半。很多人在这里踩坑浪费大量时间在版本冲突和依赖错误上。2.1 核心工具链Anaconda PyTorch我强烈推荐使用Anaconda来管理Python环境。它能为你创建一个独立的沙箱避免不同项目间的包版本互相“打架”。去Anaconda官网下载安装即可过程很简单。安装好Anaconda后打开命令行Windows用Anaconda PromptMac/Linux用终端我们创建一个专用于本项目的环境conda create -n pytorch_cnn python3.9这里指定Python 3.9因为它是一个在兼容性和稳定性上比较折中的版本。新版本如3.11、3.12可能遇到一些第三方库尚未适配的问题。激活环境conda activate pytorch_cnn接下来是安装PyTorch。这是最关键的一步版本选择错误可能导致无法使用GPU或出现各种奇怪错误。不要去PyTorch官网直接复制pip install torch torchvision这条命令官网首页的命令默认安装的是CPU版本。我们需要根据自己电脑的配置选择正确的安装命令。打开PyTorch官网pytorch.org找到“Get Started”部分你会看到一个配置选择器PyTorch Build: 选择稳定版Stable。Your OS: 选择你的操作系统。Package: 如果你用Conda管理环境就选Conda用pip也行但Conda有时能更好地处理CUDA相关的依赖。Language: Python。Compute Platform: 这是核心如果你有NVIDIA显卡并打算用GPU加速强烈推荐训练速度能快几十倍这里必须选择和你CUDA版本匹配的选项。查看CUDA版本可以在命令行输入nvidia-smi。比如如果你的CUDA版本是11.8就选择CUDA 11.8。如果没有NVIDIA显卡或者暂时不想配置CUDA就选择CPU。假设我们有CUDA 11.8选择Conda安装网站会生成如下命令conda install pytorch torchvision torchaudio pytorch-cuda11.8 -c pytorch -c nvidia在你的命令行中执行它。如果没有GPU命令会是conda install pytorch torchvision torchaudio cpuonly -c pytorch安装完成后验证一下。在Python交互环境中输入import torch print(torch.__version__) # 查看PyTorch版本 print(torch.cuda.is_available()) # 查看GPU是否可用返回True则成功如果最后一行打印True恭喜你GPU环境配置成功。如果是False请检查CUDA版本匹配和显卡驱动。注意关于AMD显卡目前PyTorch对AMD ROCm的支持不如CUDA成熟且安装过程更复杂。对于初学者如果你使用的是AMD显卡建议先使用CPU版本进行学习或者考虑在云端GPU服务器如Google Colab上实践它们通常提供现成的NVIDIA环境。2.2 辅助工具包除了PyTorch我们还需要一些帮手Torchvision: 刚才已经和PyTorch一起安装了。它提供了常用的数据集如MNIST、CIFAR-10、模型架构和图像变换工具是我们数据加载和预处理的主力。Matplotlib NumPy: 用于数据可视化和数值计算。可以用conda install matplotlib numpy安装。Jupyter Notebook/Lab: 可选但非常适合交互式开发和演示。用conda install jupyter安装。3. 理解核心组件CNN的“积木”是如何工作的在动手搭模型之前我们得先搞清楚手里的“积木”是什么。CNN的核心是卷积层但光有它还不够。3.1 卷积层Conv Layer特征提取器你可以把卷积层想象成一个拿着小窗口卷积核在图像上不断滑动巡逻的侦察兵。这个窗口在每个位置都做一次计算将窗口覆盖的局部像素值与窗口内部的权重值相乘再求和最后加上一个偏置项输出一个新的数值。这个滑动计算的过程就是“卷积”。关键参数解析in_channels: 输入数据的通道数。例如RGB彩色图像是3灰度图是1。out_channels: 你希望这一层输出多少个特征图即使用多少个不同的“侦察兵”或卷积核。每个卷积核学习检测一种特定的局部特征如边缘、纹理、颜色。kernel_size: 卷积核的尺寸。常见的有3x3, 5x5。3x3因其高效性和足够的感受野成为最主流的选择。stride: 滑动步长。默认为1即每次移动1个像素。步长为2时输出特征图的尺寸会减半。padding: 在输入图像边缘填充一圈0。这主要是为了控制输出特征图的尺寸。当padding1且kernel_size3时可以保持输出尺寸与输入尺寸相同在stride1的情况下。在PyTorch中定义一个卷积层非常简单import torch.nn as nn conv_layer nn.Conv2d(in_channels3, out_channels16, kernel_size3, stride1, padding1)这行代码创建了一个卷积层它接受3通道输入产生16个不同的特征图使用3x3的卷积核并填充1圈以保持空间尺寸。3.2 激活函数Activation Function引入非线性如果只有卷积和线性运算无论堆叠多少层整个网络本质上还是一个线性模型无法拟合复杂模式。激活函数的作用就是给网络加入非线性因素。ReLURectified Linear Unit是目前最常用的激活函数公式是f(x) max(0, x)。它把所有的负值置零正值保持不变。其优点是计算简单、能有效缓解梯度消失问题。在PyTorch中通常作为一层来使用nn.ReLU()。3.3 池化层Pooling Layer降维与不变性池化层跟在卷积和激活之后主要目的是对特征图进行下采样降低其空间尺寸宽度和高度从而减少后续层的计算量、参数数量并赋予模型一定的平移不变性即物体在图像中轻微移动输出结果不变。最大池化Max Pooling是最常用的方法。它在一个小窗口如2x2内只保留最大值。一个2x2、步长为2的最大池化层会将特征图的尺寸缩减为原来的一半。在PyTorch中定义为nn.MaxPool2d(kernel_size2, stride2)。3.4 全连接层Fully Connected Layer分类决策官经过多次“卷积-激活-池化”的循环后我们得到了高度抽象的特征图。为了最终完成分类任务需要将这些特征“拍平”Flatten成一个长向量然后输入到一个或多个全连接层中。全连接层的每个神经元都与上一层的所有神经元相连它负责综合所有提取到的特征并映射到最终的类别分数上。例如如果我们做10分类最后一个全连接层的输出维度就应该是10。3.5 网络结构设计思路一个典型的CNN结构遵循“宽深”原则随着网络加深靠近输出特征图的空间尺寸逐渐变小通过池化或大步长卷积但通道数特征图数量逐渐增加。这象征着网络从学习低级的、局部的特征如边缘、角点逐步组合成高级的、全局的、语义化的特征如眼睛、轮子、整个物体。4. 实战搭建一个用于CIFAR-10分类的CNN模型理论说再多不如动手。我们以经典的CIFAR-10数据集为例它包含6万张32x32的彩色小图片共10个类别飞机、汽车、鸟等。4.1 数据加载与预处理模型的好坏一半取决于数据。PyTorch通过torch.utils.data.DataLoader和torchvision让数据加载变得规范。import torch import torchvision import torchvision.transforms as transforms # 1. 定义数据预处理变换 transform transforms.Compose([ transforms.ToTensor(), # 将PIL图像或NumPy数组转换为PyTorch Tensor并自动缩放到[0,1] transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 对每个通道进行标准化均值0.5标准差0.5使数据范围在[-1,1]附近 ]) # 2. 下载并加载训练集和测试集 trainset torchvision.datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform) trainloader torch.utils.data.DataLoader(trainset, batch_size64, shuffleTrue, num_workers2) testset torchvision.datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtransform) testloader torch.utils.data.DataLoader(testset, batch_size64, shuffleFalse, num_workers2) # 类别名称 classes (plane, car, bird, cat, deer, dog, frog, horse, ship, truck)关键点解析ToTensor(): 必须的一步将数据转为PyTorch能处理的Tensor格式。Normalize(): 标准化能加速模型训练收敛。这里的均值和标准差是针对CIFAR-10数据集计算的经验值。batch_size: 批大小。一次迭代送入模型的数据量。太大占内存太小训练不稳定。64或128是常用起点。shuffleTrue: 仅在训练时打乱数据防止模型学习到数据顺序。num_workers: 用于数据加载的子进程数可以加快数据读取速度。根据你的CPU核心数设置。4.2 定义我们的CNN模型类现在我们来搭建模型。我们将创建一个继承自nn.Module的类。import torch.nn as nn import torch.nn.functional as F class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() # 第一个卷积块 self.conv1 nn.Conv2d(3, 32, 3, padding1) # 输入3通道输出32特征图3x3卷积填充1 self.pool nn.MaxPool2d(2, 2) # 2x2最大池化步长2 # 第二个卷积块 self.conv2 nn.Conv2d(32, 64, 3, padding1) # 输入32输出64 # 第三个卷积块 self.conv3 nn.Conv2d(64, 128, 3, padding1) # 输入64输出128 # 全连接层 # 经过三次池化32x32 - 16x16 - 8x8 - 4x4 # 最后一个卷积层输出128个4x4的特征图 self.fc1 nn.Linear(128 * 4 * 4, 256) # 拍平后输入到256维隐藏层 self.fc2 nn.Linear(256, 10) # 输出10个类别 def forward(self, x): # 前向传播定义数据流动路径 x self.pool(F.relu(self.conv1(x))) # Conv1 - ReLU - Pool x self.pool(F.relu(self.conv2(x))) # Conv2 - ReLU - Pool x self.pool(F.relu(self.conv3(x))) # Conv3 - ReLU - Pool x x.view(-1, 128 * 4 * 4) # 将特征图拍平成一个向量-1表示自动计算batch_size x F.relu(self.fc1(x)) # 全连接层1 ReLU x self.fc2(x) # 全连接层2 (输出层暂不接Softmax因损失函数包含) return x # 实例化模型 net SimpleCNN() print(net)代码走读__init__中定义了所有需要训练的参数层卷积层、全连接层。池化层没有参数但也定义在这里。forward方法定义了数据从输入到输出的完整计算图。注意F.relu是函数式调用与nn.ReLU()层效果相同。x.view(-1, 128 * 4 * 4)是形状变换的关键。-1是一个占位符会自动根据当前的batch大小进行计算。这一步将四维张量[batch, channel, height, width]转换为二维张量[batch, features]以便输入全连接层。输出层我们没有显式使用nn.Softmax因为在多分类任务中nn.CrossEntropyLoss损失函数内部已经包含了Softmax计算并且数值上更稳定。4.3 配置损失函数与优化器模型定义好后我们需要告诉它如何衡量错误损失函数以及如何根据错误来更新自己优化器。import torch.optim as optim # 定义损失函数交叉熵损失适用于多分类 criterion nn.CrossEntropyLoss() # 定义优化器随机梯度下降学习率设为0.001 optimizer optim.SGD(net.parameters(), lr0.001, momentum0.9)损失函数CrossEntropyLoss是分类任务的标准选择。它将模型的原始输出分数logits转换为概率分布并计算与真实标签的差距。优化器SGD随机梯度下降是基础优化器momentum动量参数可以帮助加速收敛并减少震荡。Adam是另一个更流行、通常需要更少调参的优化器你可以尝试optim.Adam(net.parameters(), lr0.001)。4.4 模型训练循环训练过程是一个在数据上反复迭代的循环前向传播计算预测和损失 - 反向传播计算梯度 - 优化器更新参数。# 检查是否有可用的GPU如果有则将模型和数据移至GPU device torch.device(cuda:0 if torch.cuda.is_available() else cpu) net.to(device) for epoch in range(10): # 在整个数据集上循环10次 running_loss 0.0 for i, data in enumerate(trainloader, 0): # 获取输入数据 inputs, labels data inputs, labels inputs.to(device), labels.to(device) # 移至GPU # 1. 梯度清零 optimizer.zero_grad() # 2. 前向传播 outputs net(inputs) # 3. 计算损失 loss criterion(outputs, labels) # 4. 反向传播 loss.backward() # 5. 优化器更新参数 optimizer.step() # 打印统计信息 running_loss loss.item() if i % 200 199: # 每200个mini-batch打印一次 print(f[{epoch 1}, {i 1:5d}] loss: {running_loss / 200:.3f}) running_loss 0.0 print(Finished Training)训练要点optimizer.zero_grad():至关重要PyTorch会累积梯度如果不在每次迭代前清零梯度会不断累加导致训练出错。loss.backward(): 自动计算所有参数相对于损失的梯度。optimizer.step(): 根据计算出的梯度按照优化算法如SGD更新模型参数。将模型和数据.to(device)是使用GPU加速的标准做法。代码会自动判断有GPU用GPU没有则用CPU。4.5 模型测试与评估训练完成后我们需要在从未见过的测试集上评估模型的真实能力。correct 0 total 0 # 在测试阶段不需要计算梯度用torch.no_grad()可以节省内存和计算 with torch.no_grad(): for data in testloader: images, labels data images, labels images.to(device), labels.to(device) outputs net(images) # 取输出中概率最大的类别作为预测结果 _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() print(fAccuracy of the network on the 10000 test images: {100 * correct / total:.2f} %)此外我们还可以查看模型在每个类别上的表现class_correct list(0. for i in range(10)) class_total list(0. for i in range(10)) with torch.no_grad(): for data in testloader: images, labels data images, labels images.to(device), labels.to(device) outputs net(images) _, predicted torch.max(outputs, 1) c (predicted labels).squeeze() for i in range(len(labels)): label labels[i] class_correct[label] c[i].item() class_total[label] 1 for i in range(10): if class_total[i] 0: print(fAccuracy of {classes[i]:5s}: {100 * class_correct[i] / class_total[i]:.2f} %)你可能会发现像“猫”、“狗”这类特征更复杂的类别准确率较低而“汽车”、“轮船”等特征鲜明的类别准确率较高这符合直觉。4.6 模型保存与加载训练一个好的模型可能需要很长时间我们需要保存它以便后续使用或继续训练。# 保存整个模型的结构和参数 PATH ./cifar_net.pth torch.save(net.state_dict(), PATH) # 加载模型在另一个脚本或会话中 # 首先需要重新实例化模型结构 loaded_net SimpleCNN() loaded_net.load_state_dict(torch.load(PATH)) loaded_net.to(device) loaded_net.eval() # 设置为评估模式这会关闭Dropout、BatchNorm等层的训练行为state_dict()是一个Python字典包含了模型的所有可学习参数权重和偏置。torch.save()通常保存这个字典而不是整个模型对象这样更灵活且与模型定义代码解耦。加载时必须先创建一个和保存时结构完全相同的模型实例然后再加载参数。eval()模式很重要它确保了模型在推理时行为的一致性。5. 性能优化与调试技巧实录按照上面的步骤你应该能得到一个在CIFAR-10上准确率大约在70%左右的模型。但这只是起点。下面分享一些让模型效果更好、训练更稳的实战技巧。5.1 学习率调整策略固定学习率可能不是最优的。开始时需要大步前进大学习率后期接近最优解时需要小步调优小学习率。PyTorch的torch.optim.lr_scheduler提供了多种调度器。# 使用StepLR每5个epoch将学习率乘以0.1 scheduler optim.lr_scheduler.StepLR(optimizer, step_size5, gamma0.1) # 在训练循环的每个epoch结束后调用 for epoch in range(10): # ... 训练一个epoch的代码 ... scheduler.step() # 更新学习率 print(fEpoch {epoch1}, Current LR: {optimizer.param_groups[0][lr]})5.2 防止过拟合Dropout与数据增强我们的模型可能在训练集上表现很好但在测试集上不佳这就是过拟合。对抗过拟合有两大利器Dropout在训练过程中随机“丢弃”一部分神经元将其输出置零可以防止神经元之间产生复杂的共适应关系增强模型的泛化能力。通常在全连接层之后添加。self.dropout nn.Dropout(p0.5) # 在__init__中定义 x self.dropout(F.relu(self.fc1(x))) # 在forward中使用注意在模型评估model.eval()时Dropout层会自动被关闭。数据增强Data Augmentation在训练时对输入图像进行随机变换如翻转、旋转、裁剪、颜色抖动相当于人为扩大了训练数据集。这非常有效修改训练数据的transform即可transform_train transforms.Compose([ transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.RandomCrop(32, padding4), # 随机裁剪并填充 transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) # 测试集不要做数据增强只做标准化 transform_test transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])5.3 使用TensorBoard可视化训练过程“黑箱”训练让人不安。TensorBoard可以帮助我们可视化损失曲线、准确率曲线甚至观察模型结构、参数分布和梯度流。from torch.utils.tensorboard import SummaryWriter # 在训练开始前 writer SummaryWriter(runs/cifar10_experiment_1) # 在训练循环内记录标量数据 for epoch in range(10): running_loss 0.0 for i, data in enumerate(trainloader, 0): # ... 训练步骤 ... running_loss loss.item() if i % 200 199: # 记录每200个batch的平均损失 writer.add_scalar(training loss, running_loss / 200, epoch * len(trainloader) i) running_loss 0.0 # 也可以记录每个epoch在测试集上的准确率 # ... 测试代码 ... writer.add_scalar(test accuracy, 100 * correct / total, epoch) writer.close()训练完成后在命令行运行tensorboard --logdirruns然后在浏览器打开提示的地址就能看到丰富的可视化图表了。5.4 常见错误与排查清单RuntimeError: Expected all tensors to be on the same device问题数据和模型不在同一个设备CPU/GPU上。解决确保在训练和推理时都将输入数据和模型通过.to(device)移到同一设备。RuntimeError: size mismatch, m1: [a x b], m2: [c x d]问题全连接层输入特征维度不匹配。这是新手最常见错误通常是因为卷积层输出特征图拍平后的尺寸计算错误。解决在forward函数中在view或flatten操作前打印x.shape确保计算出的特征数量与全连接层nn.Linear定义的输入维度一致。可以使用x x.view(x.size(0), -1)让PyTorch自动计算特征数但前提是你清楚经过卷积池化后的特征图尺寸。训练损失不下降Nan或保持不变可能原因1学习率设置过高或过低。尝试一个数量级的变化如0.01, 0.001, 0.0001。可能原因2忘记调用optimizer.zero_grad()导致梯度爆炸。可能原因3数据没有标准化或存在异常值。排查打印前几个batch的损失值观察梯度print(param.grad)检查输入数据范围。GPU内存溢出CUDA out of memory解决减小batch_size。这是最直接有效的方法。也可以尝试使用梯度累积技巧多次前向传播累积梯度再统一更新一次参数模拟大batch效果。验证准确率远低于训练准确率问题典型的过拟合。解决引入更强的数据增强、添加Dropout层、降低模型复杂度减少层数或通道数、使用权重衰减在优化器中设置weight_decay参数如optim.Adam(..., weight_decay1e-4)。6. 从简单CNN到现代架构下一步探索方向当你成功运行了上面的简单CNN后你已经掌握了PyTorch建模的核心流程。但这只是冰山一角。要解决更复杂的真实问题你需要了解更强大的工具和架构。6.1 使用预训练模型与迁移学习我们不需要总是从零开始训练。对于图像任务利用在ImageNet等超大数据集上预训练好的模型如ResNet, VGG, EfficientNet作为起点进行微调Fine-tuning可以极大地节省时间和计算资源并在小数据集上获得非常好的效果。Torchvision提供了这些模型的简便加载方式import torchvision.models as models # 加载预训练的ResNet18并替换最后的全连接层以适应我们的10分类任务 model_ft models.resnet18(pretrainedTrue) num_ftrs model_ft.fc.in_features # 获取原模型全连接层的输入特征数 model_ft.fc nn.Linear(num_ftrs, 10) # 替换为一个新的10分类层 # 只训练最后一层冻结前面的特征提取层 for param in model_ft.parameters(): param.requires_grad False for param in model_ft.fc.parameters(): param.requires_grad True然后你可以用很小的学习率如0.0001在CIFAR-10上训练这个新的全连接层。6.2 尝试更复杂的CNN结构我们的SimpleCNN是一个玩具模型。可以尝试搭建或使用以下经典/现代架构VGG通过堆叠多个3x3卷积来构建深度网络结构非常规整。ResNet引入了残差连接跳跃连接解决了深度网络梯度消失/爆炸的问题让训练成百上千层的网络成为可能。这是必须理解的核心思想。MobileNet使用深度可分离卷积在保持精度的同时大幅减少参数量和计算量适合移动端和嵌入式设备。6.3 引入注意力机制注意力机制Attention是当前深度学习领域的“明星”。它让模型能够动态地关注输入中更重要的部分。在CNN中你可以尝试加入通道注意力如SENet模块或空间注意力模块。例如一个简化的通道注意力模块可以这样实现class ChannelAttention(nn.Module): def __init__(self, in_channels, reduction_ratio16): super().__init__() self.avg_pool nn.AdaptiveAvgPool2d(1) self.fc nn.Sequential( nn.Linear(in_channels, in_channels // reduction_ratio, biasFalse), nn.ReLU(), nn.Linear(in_channels // reduction_ratio, in_channels, biasFalse), nn.Sigmoid() ) def forward(self, x): b, c, _, _ x.size() y self.avg_pool(x).view(b, c) y self.fc(y).view(b, c, 1, 1) return x * y.expand_as(x)将这个模块插入到你的CNN卷积块之间可以让网络自适应地重新校准每个特征通道的权重。6.4 项目扩展思路掌握了基础CNN搭建后你可以将这些知识应用到具体项目中更换数据集尝试更复杂的图像分类数据集如ImageNet的子集、你自己收集的图片。目标检测学习Faster R-CNN、YOLO等框架它们的基础仍然是CNN用于定位和识别图像中的多个物体。语义分割学习U-Net、DeepLab等模型为图像中的每个像素分类。时序数据了解1D CNN它可以用于处理文本、音频、传感器信号等序列数据。搭建第一个能跑的CNN模型就像学会了骑自行车最初可能会摇晃但一旦掌握了平衡你就打开了通往计算机视觉广阔世界的大门。最重要的是保持动手实践的习惯每学到一个新概念或新模块就立刻写代码去验证它观察它的输出和影响。PyTorch的动态性和Python的交互性为此提供了绝佳的环境。遇到错误和低准确率不要气馁那正是你深入理解系统如何工作的契机。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价