ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

PyTorch实现CNN图像分类:从数据集到模型训练与可视化

PyTorch实现CNN图像分类:从数据集到模型训练与可视化 简介面向毕业设计场景的基于卷积神经网络的图像分类系统完整项目以Python实现适合计算机、人工智能等专业学生作为课题参考或实战练手。压缩包共二十一个文件以十三个Python脚本为核心包含两个字节码文件、一组训练好的模型与数据集另有说明文档、页面模板、样式与配置文件整体体积约六十二KB内容精简但流程完整。项目同时提供TensorFlow与PyTorch两种实现每种框架下均覆盖LeNet-5、AlexNet、GoogLeNet、ResNet等经典网络便于对比学习。源码组织清晰主程序、模型定义、混淆矩阵计算与前端界面相互独立并配合类别索引文件和说明文档能够完整展示图像分类从数据准备、网络搭建、模型训练评估到可视化应用的全过程。所有源码在本地编译运行通过评审得分九十五分以上难度适中。已有二百六十九人学习使用适合需要快速搭建图像分类系统并深入理解CNN原理的开发者。1. 用 CNN 做图像分类这个毕设标题真正要解决的问题很多刚开始接触深度学习CNN图像分类毕设的人看到标题里“源码模型说明文档全部数据资料”这几个词第一反应是去找现成的包裹找到之后对着改改不动就卡住。这个标题背后的任务其实只有四件事准备图像数据集、设计一个多层卷积神经网络、把参数调到收敛、把结果可视化之后放进说明文档。读论文的时候大家都会看CNN结构图真正动手时才发现网络不收敛时首先应该排查的往往是数据和训练配置而不是模型本身。这篇文章就按毕业设计的交付顺序往下拆从环境与数据集开始把数据读取、模型搭建、训练配置、模型保存与推理串成一条可复现的链路最后一章补上答辩时更占优势的迁移学习与常见报错处理。新手能照着跑通对流程已经熟悉的人也能直接抄训练管线和参数配置。2. 搭建 CNN 图像分类模型从数据集划分到模型结构设计2.1 数据集准备与目录划分用 ImageFolder 管理图片做图像分类的第一步不是写模型而是把数据整理成 PyTorch 能直接读取的目录结构。无论是做猫狗二分类、花卉多分类还是森林图像分类最常见的交付形式就是每个类别一个文件夹。PyTorch 的ImageFolder会自动把子目录名作为类别标签省去手工制作 CSV 标签文件的麻烦。建议按以下结构放数据再做 6:2:2 的训练、验证、测试划分data/ ├── train/ │ ├── cat/ │ ├── dog/ ├── val/ │ ├── cat/ │ ├── dog/ └── test/ ├── cat/ └── dog/读取代码只需要几行from torchvision import datasets, transforms from torch.utils.data import DataLoader transform transforms.Compose([ transforms.Resize((128, 128)), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ), ]) train_dataset datasets.ImageFolder(rootdata/train, transformtransform) val_dataset datasets.ImageFolder(rootdata/val, transformtransform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers2) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers2) print(类别映射:, train_dataset.class_to_idx) print(类别列表:, train_dataset.classes)四个关键参数需要说明。Resize((128, 128))把图片统一到固定尺寸避免同一 batch 内张量形状不一致如果原图分辨率差距很大这个尺寸可以提到 224。ToTensor除了把 0~255 的像素值压缩到 0~1 之外还会把HWC排列的 NumPy 数组转成 PyTorch 需要的CHW。Normalize用 ImageNet 数据集的均值和标准差做标准化这是一组通用参数对大多数自然图像都能直接用。batch_size32是显存和内存的折中方案想提升稳定性可以试试 16想加速训练就提到 64。shuffleTrue只在训练集开启验证集保持顺序即可。2.2 写一个可解释的 CNN 结构卷积、池化、全连接怎么组合毕设里面的模型不宜太复杂越简单越容易画出结构图。下面这个三层卷积的SimpleCNN是我常用的骨架输入是 3 通道 128×128 图像输出是类别数。把它画成 CNN 结构图时每一层的张量尺寸变化一目了然import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes2): super(SimpleCNN, self).__init__() self.conv1 nn.Sequential( nn.Conv2d(in_channels3, out_channels16, kernel_size3, padding1), nn.BatchNorm2d(16), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2) ) self.conv2 nn.Sequential( nn.Conv2d(in_channels16, out_channels32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2) ) self.conv3 nn.Sequential( nn.Conv2d(in_channels32, out_channels64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2) ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(in_features64 * 16 * 16, out_features128), nn.ReLU(inplaceTrue), nn.Dropout(p0.5), nn.Linear(in_features128, out_featuresnum_classes) ) def forward(self, x): x self.conv1(x) x self.conv2(x) x self.conv3(x) x self.classifier(x) return x结构和参数要能讲清楚。输入 3×128×128经过第一个MaxPool2d变成 16×64×64第二个池化后是 32×32×32第三个池化后是 64×16×16所以全连接层的in_features是64 * 16 * 16也就是第 4 三层的输出通道数乘以宽高。kernel_size3, padding1保持卷积前后尺寸不变只增加通道数。这里没有用空洞卷积、分组卷积这类技巧因为毕设阶段把基础结构阐述清楚比堆叠深度更有利于答辩。2.3 为什么先选小模型而不是直接上最新的图像分类模型拿到毕设题目时最自然的冲动是直接套用最新的图像分类模型。但那些模型动辄几十层输入尺寸至少 224×224没有 GPU 的话训练一轮都要很久出了问题也不好定位。我一般会先用上面这个三层模型把整个流程跑通确认数据没问题、损失能降下去再考虑换更强的骨干。小模型还有一个被低估的优势训练时间和资源占用都低CPU 也能完成一次完整训练只靠 Google Colab 的免费 GPU 就能跑完 30 个 epoch。另外答辩时老师大概率会问“为什么这么设计”而不是“为什么不用最深的模型”小模型每层的作用都能用一两句话解释清楚。注意num_classes必须和ImageFolder扫描到的子目录数量一致。目录里多一个隐藏文件夹比如 macOS 的.DS_Store或者多放了一个test目录在train下会导致类别数对不上训练时不报错但预测时会出现标签索引越界。3. 训练前先确认这三个配置损失函数、优化器与数据增强参数3.1 损失与优化器CrossEntropyLoss 和 SGD 的搭配训练一个图像分类模型损失函数和优化器的选择几乎不需要纠结。分类任务默认用CrossEntropyLossPyTorch 的实现里已经包含 softmax 运算所以模型的最后一层不需要手动接nn.Softmax直接输出一个[batch_size, num_classes]的原始 logits 即可。优化器我习惯先用带动量的 SGD不用 Adam因为 SGD 对学习率的响应更规律后期调参时更容易定位问题而且论文里描述“momentum0.9”也比较规范。import torch import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleCNN(num_classes2).to(device) criterion nn.CrossEntropyLoss() optimizer optim.SGD( model.parameters(), lr0.01, momentum0.9, weight_decay1e-4 )lr0.01在这个简单模型上是一个稳妥的初始值数据量小、类别少时可以适当加大到 0.1数据噪声大就降到 0.001。momentum帮助梯度的历史方向减少震荡。weight_decay是 L2 正则项值越大模型权重被压得越紧防止过拟合的有效手段对于只有几千张图的毕设数据集1e-4到5e-4是常见区间。3.2 数据增强的三个参数翻转、裁剪、归一化训练集和验证集的预处理需要分开写。验证集只做缩放、转张量、归一化不能加随机变换否则验证指标方差很大。训练集可以加上三种常用增强方式train_transform transforms.Compose([ transforms.RandomResizedCrop(size128, scale(0.8, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ), ]) val_transform transforms.Compose([ transforms.Resize((128, 128)), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ), ])RandomResizedCrop会在原始图像上随机裁一块区域再缩放到 128×128scale(0.8, 1.0)表示裁剪面积占原图的 80% 到 100%这个参数是应对目标物体在画面中位置不固定的核心。RandomHorizontalFlip以 0.5 的概率左右翻转对猫、狗、花这类不具有方向语义的图像有效对文字识别、车牌识别这类任务必须关掉。ColorJitter的两个 0.2 表示亮度、对比度的波动的最大值数值太大会让颜色失真导致误差上升太小则增强效果不明显。3.3 训练循环与模型保存训练循环是整个项目里最不需要“创造”的部分固定的模式读取 batch、前向传播、计算损失、梯度清零、反向传播、优化器更新。但有几个细节值得抠一下。num_epochs 30 best_acc 0.0 for epoch in range(num_epochs): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss loss.item() * images.size(0) epoch_loss running_loss / len(train_dataset) model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, dim1) total labels.size(0) correct (predicted labels).sum().item() val_acc 100.0 * correct / total print(fepoch {epoch 1:02d}/{num_epochs} | loss {epoch_loss:.4f} | val acc {val_acc:.2f}%) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), model_best.pth)代码里的两个模式切换必须成对出现。model.train()让BatchNorm2d更新均值方差、让Dropout随机丢弃神经元model.eval()则固定这些层的行为。验证阶段包在torch.no_grad()里不计算梯度既省显存又加快速度。保存模型时只存state_dict()而不是整个模型好处是换环境加载时不需要保证类定义的路径完全一致只要SimpleCNN的代码在就能用load_state_dict加载。model_best.pth就是标题里“模型”目录的核心内容把准确率最高的权重单独保存避免最后 epoch 的过拟合权重污染结果。注意每个 epoch 结束后如果验证准确率下降不需要立刻干预先看完整条曲线。验证准确率前几个 epoch 不涨是正常的因为 BatchNorm 还在适应数据分布。4. 从训练到演示用模型做图像分类预测与可视化4.1 加载模型并预测单张图片top-1 与 top-5 的实现训练结束后说明文档和答辩演示里需要展示模型对新图片的预测结果。加载模型的代码和训练时几乎一样但有两个容易被忽略的参数import torch from PIL import Image model SimpleCNN(num_classes2) model.load_state_dict( torch.load(model_best.pth, map_locationcpu) ) model.eval() img Image.open(test_cat.jpg).convert(RGB) x val_transform(img).unsqueeze(0) with torch.no_grad(): logits model(x) probs torch.softmax(logits, dim1) top_p, top_idx torch.topk(probs, k2, dim1) classes [cat, dog] for i in range(top_idx.size(1)): idx top_idx[0, i].item() print(f{classes[idx]}: {top_p[0, i].item():.4f})map_locationcpu保证在没有 CUDA 的设备上也能加载 GPU 训练的权重。unsqueeze(0)把单张图像的形状从[3, 128, 128]变成[1, 3, 128, 128]因为模型要求第一个维度是 batch。torch.softmax把 logits 转成和为 1 的概率分布torch.topk返回概率最大的前 2 个类别就是常说的 top-1 和 top-5多分类时 clear 展示这个结果比只给一个最高概率更有说服力。4.2 生成混淆矩阵而不是只看准确率准确率只能回答“对了多少”回答不了“哪些类容易混”。撰写说明文档时混淆矩阵是比准确率更有价值的图表尤其当类别间存在视觉相似性时比如森林图像分类里的“林地”和“灌木丛”就很容易互相误判。from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay import matplotlib.pyplot as plt all_preds [] all_labels [] model.eval() with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, preds torch.max(outputs, dim1) all_preds.extend(preds.cpu().tolist()) all_labels.extend(labels.tolist()) cm confusion_matrix(all_labels, all_preds) disp ConfusionMatrixDisplay(confusion_matrixcm, display_labelstrain_dataset.classes) disp.plot(cmapBlues) plt.savefig(confusion_matrix.png, dpi150)验证集按类别分布不均匀时准确率可能虚高。如果“猫”占了验证集 90% 的样本模型即使把“狗”全部判错准确率也还有 90%。混淆矩阵可以直接看到这类问题在答辩时遇到“你的准确率为什么这么高”的提问用这张图能给出最直观的解释。4.3 把训练曲线画出来验证 epoch 参数是否合理训练代码里已经保存了每个 epoch 的损失和验证准确率把它们收集到列表里画成折线图是判断训练是否收敛的关键依据import matplotlib.pyplot as plt fig, ax1 plt.subplots() ax1.plot(range(1, num_epochs 1), train_losses, labeltrain loss, colortab:red) ax1.set_xlabel(epoch) ax1.set_ylabel(loss) ax2 ax1.twinx() ax2.plot(range(1, num_epochs 1), val_accs, labelval acc, colortab:blue) ax2.set_ylabel(accuracy (%)) plt.savefig(training_curve.png, dpi150)画双曲线图是把两个量纲不同的指标放在同一个坐标轴里的标准做法。如果train_loss持续下降但val_acc不涨甚至下降就是典型的过拟合信号需要回到第 3 章调数据增强或weight_decay。如果两个曲线都在震荡说明学习率偏大把lr降到0.001重新训练。5. 在毕设答辩中加分迁移学习、骨干网络替换与常见报错处理5.1 从自己写的 CNN 换成预训练模型小模型跑通后想进一步提升准确率或让项目看起来更有深度最常见且可靠的做法是迁移学习。PyTorch 里替换骨干网络只需几行代码import torchvision.models as models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_features model.fc.in_features model.fc nn.Linear(num_features, num_classes) # 冻结除最后全连接层以外的全部参数 for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad Truerequires_gradFalse的含义是反向传播时不对这些参数计算梯度训练时只更新最后一层。冻结骨干网络可以大幅减少计算量和显存占用batch_size 可以保持 32 不变训练速度也不慢。如果希望效果更好可以解冻最后几个残差块一起微调但学习率要降到0.0001这个量级否则很容易把预训练权重破坏掉。5.2 小数据集过拟合的三个应对顺序当训练变慢或验证集表现差时按这个顺序排查先减少 epoch 数量或用早停法观察验证精度停止上升就回头加载model_best.pth其次调大RandomResizedCrop的裁剪范围并加一点ColorJitter数据多样性上去了过拟合自然缓解最后再调整weight_decay到1e-3。不要一开始就换模型骨干网络替换应排在这些数据层面的手段之后。5.3 三个常见报错与修正方法报错信息原因修正方式shape [] is invalid for input of size 0某个类别的文件夹下没有图片检查数据集目录删除空文件夹Target 2 is out of bounds标签索引超过num_classesnum_classes改为与ImageFolder扫描到的目录数一致CUDA out of memorybatch_size 太大或图片尺寸太大batch_size 降到 16 或 8图片尺寸降到 96报错信息本身不是重点定位思路才是。这一类错误在 PyTorch 里几乎都能通过打印len(train_dataset.classes)和train_dataset.class_to_idx快速发现写说明文档时也可以把这句打印输出作为数据集验证的一部分写进去让审查者确信你的数据预处理流程是可复现的。最后再补一个实用技巧把所有超参数包括学习率、batch_size、epoch、图片尺寸、类别数集中写到一个config.py里训练代码和说明文档里的参数表都从这份配置引用。答辩时被问到“参数怎么设的”就可以直接展示这份文件比翻训练代码里的散落数值更有说服力。本文还有配套的精品资源点击获取
返回列表