简介面向计算机相关专业毕业设计及课程实践场景提供基于Pytorch的CNN垃圾分类系统完整项目源码。项目经导师指导认可评审分98分所有代码均已本地编译并调试通过难度适中适合毕业设计、课程设计、期末大作业或项目实战练习可帮助学习者快速掌握图像分类模型从搭建到部署的完整流程。压缩包内共2000个文件包含126个Python源码文件、1866张jpg样本图片、5个xml标注文件及3个txt说明文件整体大小约285MB。jpg图片为垃圾分类样本数据集可直接用于模型训练与验证Python源码覆盖数据加载、模型构建、训练评估等核心模块结构清晰便于二次开发。目前已有171人学习下载具备一定参考价值。通过该项目可深入理解CNN网络结构、数据集预处理与分类任务实现思路并可将代码迁移到其他图像识别场景是较完整的深度学习实践范例。1. 垃圾分类为什么选择CNN而不是传统视觉方案拿到这套PyTorch的CNN垃圾分类系统源码我第一反应是看数据组织方式——项目没有花哨目录树只有trash73.jpg、trash95.jpg、trash34.jpg这一批按编号命名的图片类别信息需要再加工才能变成训练集。这种命名在毕设里很常见但直接丢给训练脚本一定会出错。这个高分项目的核心价值是用CNN卷积神经网络解决真实的小型图像分类问题输入一张垃圾照片输出类别和置信度。网络不复杂但数据管线、训练策略、评估脚本完整在Python与PyTorch环境里可直接运行。适合正在做毕业设计的计算机专业学生也适合想用简短代码理解CNN训练闭环的开发者。下面按数据加载、模型搭建、训练优化、推理验证四条线拆开讲。2. 数据集划分与DataLoader管线的PyTorch实现2.1 从文件名到结构化数据集的划分策略图片文件名里带的是编号不是类别比如trash73.jpg、trash119.jpg、trash48.jpg。要在CNN监督训练里用起来第一步是把这些扁平文件映射成(图片路径, 类别序号)配对。我不建议在读取时临时解析文件名做规则匹配那样一旦命名规则变化整个训练流程就要跟着改。更稳妥的做法是先做一次目录归档每个类别一个子目录再按比例拆出train、val两份。我一般用8:2划分。垃圾分类场景下图片总量通常在几千张以内验证集给得太多训练集不够学纹理给太少val指标波动几下就把最优模型错误丢弃。如果原始图片按trash编号平铺则先抽一定比例图片做验证集剩余再按类别归档确保验证集中不会混入训练阶段见过的同类样本。划分脚本运行前固定random.seed与numpy.random.seed保证同一份源码在不同机器上跑出的训练集、验证集划分一致。答辩时评委只要看到两次运行结果可复现就会认为整个实验流程是规范的。选Anaconda做Python环境的话用conda create -n trash python3.8建虚拟环境再按GPU驱动版本安装对应PyTorch这套源码在CPU或CUDA两种环境下都能跑。train/ recyclable/ trash73.jpg kitchen/ trash95.jpg harmful/ trash34.jpg other/ trash48.jpg val/ recyclable/ trash119.jpg ...类别子目录名建议直接用英文不带空格。中文目录在Windows下配合torchvision.datasets.ImageFolder可能出现编码报错而自定义Dataset读取时用英文路径也能避免UnicodeDecodeError干扰训练主流程。2.2 Dataset子类实现与图片解码细节数据归档完成后写一个继承自torch.utils.data.Dataset的TrashDataset。直接用torchvision.datasets.ImageFolder可以少写代码但很多毕设项目源码习惯自定义后续要挂特殊采样器或做类别加权时更好改。这里给出的是项目里最常见的一种实现。import os import torch from PIL import Image from torch.utils.data import Dataset class TrashDataset(Dataset): 垃圾分类数据集封装读取目录结构为 class_name/*.jpg 的样本。 def __init__(self, root_dir, transformNone): self.image_paths [] self.labels [] self.transform transform # 子目录名按字母排序后映射为类别索引保证多次运行结果一致 self.classes sorted(os.listdir(root_dir)) self.class_to_idx {cls_name: i for i, cls_name in enumerate(self.classes)} for cls_name in self.classes: cls_dir os.path.join(root_dir, cls_name) for fname in os.listdir(cls_dir): if fname.lower().endswith((.jpg, .jpeg, .png)): self.image_paths.append(os.path.join(cls_dir, fname)) self.labels.append(self.class_to_idx[cls_name]) def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img Image.open(self.image_paths[idx]).convert(RGB) label self.labels[idx] if self.transform: img self.transform(img) return img, label这里三个细节值得注意。convert(RGB)把灰度图或带alpha通道的PNG统一转成三通道防止单张图片通道数不一致导致后续torch.stack维度冲突。sorted(os.listdir(root_dir))保证类别编号在一次训练流程内稳定训练中断后重启模型权重才不至于错位。fname.lower().endswith()把.jpg、.jpeg、.png都纳入避免不同来源的垃圾图片因扩展名大小写问题被跳过。PIL在样本数量少时解码不是瓶颈但如果本机跑过原始图片缓存建议先把图片统一缩到256x256再存入本地缓存目录。训练循环里读取小图IO压力小很多。需要注意缓存目录与原始目录分开否则Dataset构造时会递归扫到重复文件。2.3 数据增强策略与transforms参数设计垃圾照片的拍摄条件比ImageNet复杂光线不均匀、背景杂物、瓶子横竖摆放不一致。数据增强是在不增加样本数的情况下扩展数据分布这部分参数直接影响最终准确率。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.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((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])验证集不加随机增强只做确定性Resize和Normalize。若把翻转或旋转加到验证集验证准确率会随随机种子抖动早停判断失真。下面是这套源码调参时的一张对照表。参数作用建议值设太大的后果RandomHorizontalFlip镜像样本增强对称性0.5带文字的可回收瓶分类失效RandomRotation模拟摆放角度偏移15度超过30度时黑边干扰特征提取ColorJitter亮度/对比度/饱和度扰动0.2真实颜色信息被破坏Resize后是否Crop模拟局部观察视角不裁剪类别区域被裁掉影响识别需要说明的是RandomResizedCrop在垃圾分类里不一定要用。垃圾图像类别判别依赖的往往是瓶身整体轮廓裁剪后局部信息不足以支撑分类。项目源码如果直接统一Resize到224x224说明作者优先保证形状稳定性这在小数据集上是合理选择。2.4 DataLoader参数设定与多进程加载的坑Dataset定义好之后用DataLoader包装。注意num_workers是毕设电脑上最敏感的参数。from torch.utils.data import DataLoader train_loader DataLoader( datasettrain_dataset, batch_size32, shuffleTrue, num_workers0, pin_memoryTrue, drop_lastTrue ) val_loader DataLoader( datasetval_dataset, batch_size32, shuffleFalse, num_workers0, pin_memoryTrue )num_workers0表示在主进程中加载数据不会触发Windows上spawn模式的多进程问题。很多教程建议把num_workers调成4或8但如果在Windows上直接跑频繁报“DataLoader worker exited unexpectedly”最可能就是num_workers大于0导致子进程重新导入Dataset时递归异常。毕设项目能用CPU训练就不必纠结这个参数。drop_lastTrue丢弃最后一个不完整batch避免BatchNorm在batch size为1时统计量漂移。提示验证集Loader保持shuffleFalse保证每个epoch遍历顺序一致混淆矩阵和逐类准确率才有对比意义。3. CNN网络结构搭建Conv、BN、Pooling的参数设计3.1 为什么用自定义轻量CNN而不是直接加载ResNet毕设项目的常规偷懒路线是torchvision里直接resnet18(pretrainedTrue)替换最后一层全连接二十行代码就能训练。这套源码没有这样处理而是从零搭了一个CNN原因有两层。第一垃圾分类作为小样本任务类间差异集中在瓶身形状、材质纹理、颜色上。深层预训练网络在ImageNet上学到的特征偏通用直接迁移到垃圾图片后底层边缘特征虽然有用但高层语义特征与垃圾类别无关微调收益有限。第二参数量差距带来的训练开销很大。ResNet18单次前向计算约1.8 GFLOPs而一个3卷积块的轻量CNN只有约0.3 GFLOPs训练单epoch时间相差近5倍。在纯CPU环境下跑毕设项目这个差距会直接决定能否在截止日期前调完参数。3.2 自定义CNN结构代码import torch.nn as nn class TrashCNN(nn.Module): def __init__(self, num_classes4): super(TrashCNN, self).__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(128, 128), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(128, num_classes) ) def forward(self, x): return self.classifier(self.features(x))整个网络由3个卷积块加2层全连接组成。卷积块统一用kernel_size3, padding1保证特征图宽高在卷积前后不变只由MaxPool2d做2倍降采样。图片从224x224经过三次池化变成28x28通道数从3增至128。AdaptiveAvgPool2d((1,1))把任意尺寸的特征图压缩为1x1彻底解除全连接层对输入尺寸的依赖。BN层放在卷积层与ReLU之间。批归一化让每层输入分布稳定训练时可以适当提高学习率。若去掉BN小数据集上loss下降速度明显变慢最终准确率通常也要低3到5个百分点。Dropout只在全连接部分用卷积层不设Dropout因为卷积层参数共享特性已经自带正则效果。3.3 各层输出尺寸与参数量对照搭建CNN时必问的答辩问题是“每层输出多大、参数量多少”。下面这张表可以直接用。层名输入尺寸输出尺寸参数量Conv13x224x22432x224x224896MaxPool132x224x22432x112x1120Conv232x112x11264x112x11218496MaxPool264x112x11264x56x560Conv364x56x56128x56x5673856MaxPool3128x56x56128x28x280Linear112812816512Linear21284516卷积层参数量按C_out x C_in x K x K计算Linear1是128x128权重加128偏置。整个模型参数量约11万权重文件不到0.5MBCPU推理单张图几十毫秒。这个尺寸对毕设展示和移动端部署都很友好也是源码被评为高分的一个技术点。3.4 forward张量形状校验方法训练中遇到维度报错不要反复跑训练脚本。用一段极简代码打印每一层输出形状十秒内定位问题。import torch model TrashCNN() fake_input torch.randn(1, 3, 224, 224) for name, layer in model.features.named_children(): fake_input layer(fake_input) print(name, fake_input.shape)输出结果是Conv不改变宽高MaxPool把宽高减半。自定义CNN最容易翻车的是Flatten后的维度如果不用AdaptiveAvgPool2d而直接算出特征图尺寸然后手动拼Linear换数据集尺寸后全连接维度就崩了。用AdaptiveAvgPool2d((1,1))后224、256、320任意输入尺寸最终都变成1x1x128全连接层完全不需要改。4. 训练流程、优化器选择与早停机制4.1 损失函数与优化器组合多分类任务默认用CrossEntropyLoss。PyTorch的CrossEntropyLoss内部把LogSoftmax和NLLLoss合并成一个算子模型最后一层输出logits即可不要手动加Softmax。若在分类器输出层额外塞Softmax再算loss梯度会不稳定训练loss下降变慢。优化器这里推荐AdamW。相比原始AdamAdamW把weight decay从梯度修正中解耦正则效果更干净。学习率取1e-3weight_decay取1e-4。SGD with momentum不是不行但需要更长的预热步数配合余弦退火在课设毕设时间紧张的情况下不划算。实际对比结果自定义CNN在AdamW下第10轮左右验证集基本收敛而SGD往往到第20轮还在震荡。4.2 训练循环完整实现训练与验证放在同一个epoch循环里每个epoch末尾做一次验证同时维护一个best_acc变量存储历史最优权重。import torch import torch.nn as nn from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR device torch.device(cuda if torch.cuda.is_available() else cpu) model TrashCNN(num_classes4).to(device) criterion nn.CrossEntropyLoss() optimizer AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max30, eta_min1e-5) best_acc 0.0 num_epochs 30 for epoch in range(num_epochs): model.train() total_loss 0.0 correct 0 total 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() # 梯度裁剪防止个别异常样本导致梯度爆炸 nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) optimizer.step() total_loss loss.item() * images.size(0) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) scheduler.step() train_acc correct / total # 验证阶段 model.eval() val_correct 0 val_total 0 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, 1) val_correct (preds labels).sum().item() val_total labels.size(0) val_acc val_correct / val_total if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(fEpoch {epoch1}/{num_epochs} | fLoss {total_loss/total:.4f} | fTrain Acc {train_acc:.4f} | fVal Acc {val_acc:.4f})验证阶段必须用torch.no_grad()包裹否则PyTorch会为验证分支建计算图显存占用翻倍。梯度裁剪的max_norm5.0在CNN小网络上不会频繁触发但能防止某张异常图片把BN的running_mean带偏。best_model.pth存的是state_dict而非完整模型加载时需要先实例化模型结构再执行load_state_dict。4.3 学习率调度与超参对照参数默认值调优方向batch_size32显存不足降到16同步观察BN统计量learning_rate1e-3loss不降时降到1e-4weight_decay1e-4过拟合时升到5e-4T_max30与总epoch对齐max_norm5.0梯度数值异常时降到2.0如果是CPU训练batch_size建议保持32以下。神经网络在CPU上的瓶颈通常是卷积计算batch太大反而导致内存交换频繁单epoch耗时线性甚至超线性增长。在源码的默认配置下CPU上每个epoch大约需要2到3分钟30轮训练约一个半小时适合上午跑任务、下午做分析。4.4 训练日志排查技巧训练日志里三列指标要放在一起看。train_loss稳步下降但val_acc不涨说明模型已经记住训练集特征当前增强力度不够优先调ColorJitter或Rotation参数而不是增加模型层数。第5轮前train_acc冲到0.95以上而val_acc卡在0.6最可能是数据集的图片归档错误——去训练目录里抽查每个类别文件夹看看是否有图片被放错位置。还有一类特殊现象是train_loss先降后升。原因是学习率太大导致更新步长跨过了最优点交叉熵在logits接近正确类别时梯度方向突变。把学习率减半或者把eta_min设到1e-6曲线就会回复平缓。区分是学习率问题还是数据问题最简单的办法是固定随机种子重复一次训练如果曲线形状完全相同基本可排除数据混叠聚焦到优化器参数上。5. 单张图片推理与混淆矩阵验证5.1 单图推理脚本与batch维度处理模型训练好最终要落到的场景是对一张新图给出类别和置信度。推理代码与训练循环很像三个关键差异是关闭梯度、输入增加batch维、不需要真实标签。import torch from PIL import Image import torchvision.transforms as transforms def predict_single(model, image_path, class_names, devicecpu): # 推理用的transform必须与验证时保持一致 transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.open(image_path).convert(RGB) x transform(img).unsqueeze(0) # (3,224,224) - (1,3,224,224) x x.to(device) model.eval() with torch.no_grad(): logits model(x) probs torch.softmax(logits, dim1) # 转概率 prob, idx torch.max(probs, dim1) return class_names[idx.item()], prob.item()unsqueeze(0)是最容易被忽略的地方。模型forward期望的是4D张量(N, C, H, W)单张图片读取后只有3D尺寸必须补上batch维度。最终softmax在dim1上进行得到每个类别的概率分布后取最大值对应的类这就是置信度最高的预测。源码项目里如果配了gradio或flask接口最终调用的也是这个函数主体。5.2 混淆矩阵绘制与错误样本定位单纯看val_acc会掩盖类别之间的差异。分批次的垃圾类别中可回收物和有害垃圾在塑料材质外观上相近准确率往往最低。用sklearn的confusion_matrix把预测与真实标签画出来第一时间就能看到哪些类别互相混淆。import numpy as np from sklearn.metrics import confusion_matrix, classification_report import seaborn as sns import matplotlib.pyplot as plt model.eval() all_preds [] all_labels [] with torch.no_grad(): for images, labels in val_loader: images images.to(device) outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) print(classification_report(all_labels, all_preds)) # 画图时把类别名标上 sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelstrain_dataset.classes, yticklabelstrain_dataset.classes) plt.show()混淆矩阵的对角线值越高说明每个类别的特征分得越开。查看被错分的样本图时重点是看图片的背景信息——垃圾图像里大量样本的背景是桌面、地面、手部如果某一类样本背景都集中在同一种场景CNN很容易把背景特征当作类别特征。处理办法是增强里加RandomResizedCrop(scale(0.6, 1.0))强制模型关注目标区域而不是背景这也是在原始训练参数基础上的有效进阶调整。5.3 样本量不足时冻结前层的小技巧当某个类别图片只有几十张、训练准确率始终上不去时一个实用技巧是冻结第一个卷积块只优化后面的层。将requires_grad置为False后Adam只更新后半部分参数可学习参数量直接缩减约10%过拟合速度明显变慢。for name, param in model.named_parameters(): if name.startswith(features.0): param.requires_grad False冻结后第一个卷积块仍然参与前向计算但梯度不再回传。这套做法适合小样本迁移场景冻结浅层的边缘、颜色特征提取能力同时释放深层的类别判别能力在垃圾分类这种细粒度不高的任务上收益明显。要快速验证模型加载是否正常用torch.save存权重后执行一次加载比对验证集的准确率是否有差异即可确认推理链路没有断点。本文还有配套的精品资源点击获取
PyTorch实现CNN垃圾分类：从数据处理到模型推理
简介面向计算机相关专业毕业设计及课程实践场景提供基于Pytorch的CNN垃圾分类系统完整项目源码。项目经导师指导认可评审分98分所有代码均已本地编译并调试通过难度适中适合毕业设计、课程设计、期末大作业或项目实战练习可帮助学习者快速掌握图像分类模型从搭建到部署的完整流程。压缩包内共2000个文件包含126个Python源码文件、1866张jpg样本图片、5个xml标注文件及3个txt说明文件整体大小约285MB。jpg图片为垃圾分类样本数据集可直接用于模型训练与验证Python源码覆盖数据加载、模型构建、训练评估等核心模块结构清晰便于二次开发。目前已有171人学习下载具备一定参考价值。通过该项目可深入理解CNN网络结构、数据集预处理与分类任务实现思路并可将代码迁移到其他图像识别场景是较完整的深度学习实践范例。1. 垃圾分类为什么选择CNN而不是传统视觉方案拿到这套PyTorch的CNN垃圾分类系统源码我第一反应是看数据组织方式——项目没有花哨目录树只有trash73.jpg、trash95.jpg、trash34.jpg这一批按编号命名的图片类别信息需要再加工才能变成训练集。这种命名在毕设里很常见但直接丢给训练脚本一定会出错。这个高分项目的核心价值是用CNN卷积神经网络解决真实的小型图像分类问题输入一张垃圾照片输出类别和置信度。网络不复杂但数据管线、训练策略、评估脚本完整在Python与PyTorch环境里可直接运行。适合正在做毕业设计的计算机专业学生也适合想用简短代码理解CNN训练闭环的开发者。下面按数据加载、模型搭建、训练优化、推理验证四条线拆开讲。2. 数据集划分与DataLoader管线的PyTorch实现2.1 从文件名到结构化数据集的划分策略图片文件名里带的是编号不是类别比如trash73.jpg、trash119.jpg、trash48.jpg。要在CNN监督训练里用起来第一步是把这些扁平文件映射成(图片路径, 类别序号)配对。我不建议在读取时临时解析文件名做规则匹配那样一旦命名规则变化整个训练流程就要跟着改。更稳妥的做法是先做一次目录归档每个类别一个子目录再按比例拆出train、val两份。我一般用8:2划分。垃圾分类场景下图片总量通常在几千张以内验证集给得太多训练集不够学纹理给太少val指标波动几下就把最优模型错误丢弃。如果原始图片按trash编号平铺则先抽一定比例图片做验证集剩余再按类别归档确保验证集中不会混入训练阶段见过的同类样本。划分脚本运行前固定random.seed与numpy.random.seed保证同一份源码在不同机器上跑出的训练集、验证集划分一致。答辩时评委只要看到两次运行结果可复现就会认为整个实验流程是规范的。选Anaconda做Python环境的话用conda create -n trash python3.8建虚拟环境再按GPU驱动版本安装对应PyTorch这套源码在CPU或CUDA两种环境下都能跑。train/ recyclable/ trash73.jpg kitchen/ trash95.jpg harmful/ trash34.jpg other/ trash48.jpg val/ recyclable/ trash119.jpg ...类别子目录名建议直接用英文不带空格。中文目录在Windows下配合torchvision.datasets.ImageFolder可能出现编码报错而自定义Dataset读取时用英文路径也能避免UnicodeDecodeError干扰训练主流程。2.2 Dataset子类实现与图片解码细节数据归档完成后写一个继承自torch.utils.data.Dataset的TrashDataset。直接用torchvision.datasets.ImageFolder可以少写代码但很多毕设项目源码习惯自定义后续要挂特殊采样器或做类别加权时更好改。这里给出的是项目里最常见的一种实现。import os import torch from PIL import Image from torch.utils.data import Dataset class TrashDataset(Dataset): 垃圾分类数据集封装读取目录结构为 class_name/*.jpg 的样本。 def __init__(self, root_dir, transformNone): self.image_paths [] self.labels [] self.transform transform # 子目录名按字母排序后映射为类别索引保证多次运行结果一致 self.classes sorted(os.listdir(root_dir)) self.class_to_idx {cls_name: i for i, cls_name in enumerate(self.classes)} for cls_name in self.classes: cls_dir os.path.join(root_dir, cls_name) for fname in os.listdir(cls_dir): if fname.lower().endswith((.jpg, .jpeg, .png)): self.image_paths.append(os.path.join(cls_dir, fname)) self.labels.append(self.class_to_idx[cls_name]) def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img Image.open(self.image_paths[idx]).convert(RGB) label self.labels[idx] if self.transform: img self.transform(img) return img, label这里三个细节值得注意。convert(RGB)把灰度图或带alpha通道的PNG统一转成三通道防止单张图片通道数不一致导致后续torch.stack维度冲突。sorted(os.listdir(root_dir))保证类别编号在一次训练流程内稳定训练中断后重启模型权重才不至于错位。fname.lower().endswith()把.jpg、.jpeg、.png都纳入避免不同来源的垃圾图片因扩展名大小写问题被跳过。PIL在样本数量少时解码不是瓶颈但如果本机跑过原始图片缓存建议先把图片统一缩到256x256再存入本地缓存目录。训练循环里读取小图IO压力小很多。需要注意缓存目录与原始目录分开否则Dataset构造时会递归扫到重复文件。2.3 数据增强策略与transforms参数设计垃圾照片的拍摄条件比ImageNet复杂光线不均匀、背景杂物、瓶子横竖摆放不一致。数据增强是在不增加样本数的情况下扩展数据分布这部分参数直接影响最终准确率。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.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((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])验证集不加随机增强只做确定性Resize和Normalize。若把翻转或旋转加到验证集验证准确率会随随机种子抖动早停判断失真。下面是这套源码调参时的一张对照表。参数作用建议值设太大的后果RandomHorizontalFlip镜像样本增强对称性0.5带文字的可回收瓶分类失效RandomRotation模拟摆放角度偏移15度超过30度时黑边干扰特征提取ColorJitter亮度/对比度/饱和度扰动0.2真实颜色信息被破坏Resize后是否Crop模拟局部观察视角不裁剪类别区域被裁掉影响识别需要说明的是RandomResizedCrop在垃圾分类里不一定要用。垃圾图像类别判别依赖的往往是瓶身整体轮廓裁剪后局部信息不足以支撑分类。项目源码如果直接统一Resize到224x224说明作者优先保证形状稳定性这在小数据集上是合理选择。2.4 DataLoader参数设定与多进程加载的坑Dataset定义好之后用DataLoader包装。注意num_workers是毕设电脑上最敏感的参数。from torch.utils.data import DataLoader train_loader DataLoader( datasettrain_dataset, batch_size32, shuffleTrue, num_workers0, pin_memoryTrue, drop_lastTrue ) val_loader DataLoader( datasetval_dataset, batch_size32, shuffleFalse, num_workers0, pin_memoryTrue )num_workers0表示在主进程中加载数据不会触发Windows上spawn模式的多进程问题。很多教程建议把num_workers调成4或8但如果在Windows上直接跑频繁报“DataLoader worker exited unexpectedly”最可能就是num_workers大于0导致子进程重新导入Dataset时递归异常。毕设项目能用CPU训练就不必纠结这个参数。drop_lastTrue丢弃最后一个不完整batch避免BatchNorm在batch size为1时统计量漂移。提示验证集Loader保持shuffleFalse保证每个epoch遍历顺序一致混淆矩阵和逐类准确率才有对比意义。3. CNN网络结构搭建Conv、BN、Pooling的参数设计3.1 为什么用自定义轻量CNN而不是直接加载ResNet毕设项目的常规偷懒路线是torchvision里直接resnet18(pretrainedTrue)替换最后一层全连接二十行代码就能训练。这套源码没有这样处理而是从零搭了一个CNN原因有两层。第一垃圾分类作为小样本任务类间差异集中在瓶身形状、材质纹理、颜色上。深层预训练网络在ImageNet上学到的特征偏通用直接迁移到垃圾图片后底层边缘特征虽然有用但高层语义特征与垃圾类别无关微调收益有限。第二参数量差距带来的训练开销很大。ResNet18单次前向计算约1.8 GFLOPs而一个3卷积块的轻量CNN只有约0.3 GFLOPs训练单epoch时间相差近5倍。在纯CPU环境下跑毕设项目这个差距会直接决定能否在截止日期前调完参数。3.2 自定义CNN结构代码import torch.nn as nn class TrashCNN(nn.Module): def __init__(self, num_classes4): super(TrashCNN, self).__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(128, 128), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(128, num_classes) ) def forward(self, x): return self.classifier(self.features(x))整个网络由3个卷积块加2层全连接组成。卷积块统一用kernel_size3, padding1保证特征图宽高在卷积前后不变只由MaxPool2d做2倍降采样。图片从224x224经过三次池化变成28x28通道数从3增至128。AdaptiveAvgPool2d((1,1))把任意尺寸的特征图压缩为1x1彻底解除全连接层对输入尺寸的依赖。BN层放在卷积层与ReLU之间。批归一化让每层输入分布稳定训练时可以适当提高学习率。若去掉BN小数据集上loss下降速度明显变慢最终准确率通常也要低3到5个百分点。Dropout只在全连接部分用卷积层不设Dropout因为卷积层参数共享特性已经自带正则效果。3.3 各层输出尺寸与参数量对照搭建CNN时必问的答辩问题是“每层输出多大、参数量多少”。下面这张表可以直接用。层名输入尺寸输出尺寸参数量Conv13x224x22432x224x224896MaxPool132x224x22432x112x1120Conv232x112x11264x112x11218496MaxPool264x112x11264x56x560Conv364x56x56128x56x5673856MaxPool3128x56x56128x28x280Linear112812816512Linear21284516卷积层参数量按C_out x C_in x K x K计算Linear1是128x128权重加128偏置。整个模型参数量约11万权重文件不到0.5MBCPU推理单张图几十毫秒。这个尺寸对毕设展示和移动端部署都很友好也是源码被评为高分的一个技术点。3.4 forward张量形状校验方法训练中遇到维度报错不要反复跑训练脚本。用一段极简代码打印每一层输出形状十秒内定位问题。import torch model TrashCNN() fake_input torch.randn(1, 3, 224, 224) for name, layer in model.features.named_children(): fake_input layer(fake_input) print(name, fake_input.shape)输出结果是Conv不改变宽高MaxPool把宽高减半。自定义CNN最容易翻车的是Flatten后的维度如果不用AdaptiveAvgPool2d而直接算出特征图尺寸然后手动拼Linear换数据集尺寸后全连接维度就崩了。用AdaptiveAvgPool2d((1,1))后224、256、320任意输入尺寸最终都变成1x1x128全连接层完全不需要改。4. 训练流程、优化器选择与早停机制4.1 损失函数与优化器组合多分类任务默认用CrossEntropyLoss。PyTorch的CrossEntropyLoss内部把LogSoftmax和NLLLoss合并成一个算子模型最后一层输出logits即可不要手动加Softmax。若在分类器输出层额外塞Softmax再算loss梯度会不稳定训练loss下降变慢。优化器这里推荐AdamW。相比原始AdamAdamW把weight decay从梯度修正中解耦正则效果更干净。学习率取1e-3weight_decay取1e-4。SGD with momentum不是不行但需要更长的预热步数配合余弦退火在课设毕设时间紧张的情况下不划算。实际对比结果自定义CNN在AdamW下第10轮左右验证集基本收敛而SGD往往到第20轮还在震荡。4.2 训练循环完整实现训练与验证放在同一个epoch循环里每个epoch末尾做一次验证同时维护一个best_acc变量存储历史最优权重。import torch import torch.nn as nn from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR device torch.device(cuda if torch.cuda.is_available() else cpu) model TrashCNN(num_classes4).to(device) criterion nn.CrossEntropyLoss() optimizer AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max30, eta_min1e-5) best_acc 0.0 num_epochs 30 for epoch in range(num_epochs): model.train() total_loss 0.0 correct 0 total 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() # 梯度裁剪防止个别异常样本导致梯度爆炸 nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) optimizer.step() total_loss loss.item() * images.size(0) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) scheduler.step() train_acc correct / total # 验证阶段 model.eval() val_correct 0 val_total 0 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, 1) val_correct (preds labels).sum().item() val_total labels.size(0) val_acc val_correct / val_total if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(fEpoch {epoch1}/{num_epochs} | fLoss {total_loss/total:.4f} | fTrain Acc {train_acc:.4f} | fVal Acc {val_acc:.4f})验证阶段必须用torch.no_grad()包裹否则PyTorch会为验证分支建计算图显存占用翻倍。梯度裁剪的max_norm5.0在CNN小网络上不会频繁触发但能防止某张异常图片把BN的running_mean带偏。best_model.pth存的是state_dict而非完整模型加载时需要先实例化模型结构再执行load_state_dict。4.3 学习率调度与超参对照参数默认值调优方向batch_size32显存不足降到16同步观察BN统计量learning_rate1e-3loss不降时降到1e-4weight_decay1e-4过拟合时升到5e-4T_max30与总epoch对齐max_norm5.0梯度数值异常时降到2.0如果是CPU训练batch_size建议保持32以下。神经网络在CPU上的瓶颈通常是卷积计算batch太大反而导致内存交换频繁单epoch耗时线性甚至超线性增长。在源码的默认配置下CPU上每个epoch大约需要2到3分钟30轮训练约一个半小时适合上午跑任务、下午做分析。4.4 训练日志排查技巧训练日志里三列指标要放在一起看。train_loss稳步下降但val_acc不涨说明模型已经记住训练集特征当前增强力度不够优先调ColorJitter或Rotation参数而不是增加模型层数。第5轮前train_acc冲到0.95以上而val_acc卡在0.6最可能是数据集的图片归档错误——去训练目录里抽查每个类别文件夹看看是否有图片被放错位置。还有一类特殊现象是train_loss先降后升。原因是学习率太大导致更新步长跨过了最优点交叉熵在logits接近正确类别时梯度方向突变。把学习率减半或者把eta_min设到1e-6曲线就会回复平缓。区分是学习率问题还是数据问题最简单的办法是固定随机种子重复一次训练如果曲线形状完全相同基本可排除数据混叠聚焦到优化器参数上。5. 单张图片推理与混淆矩阵验证5.1 单图推理脚本与batch维度处理模型训练好最终要落到的场景是对一张新图给出类别和置信度。推理代码与训练循环很像三个关键差异是关闭梯度、输入增加batch维、不需要真实标签。import torch from PIL import Image import torchvision.transforms as transforms def predict_single(model, image_path, class_names, devicecpu): # 推理用的transform必须与验证时保持一致 transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.open(image_path).convert(RGB) x transform(img).unsqueeze(0) # (3,224,224) - (1,3,224,224) x x.to(device) model.eval() with torch.no_grad(): logits model(x) probs torch.softmax(logits, dim1) # 转概率 prob, idx torch.max(probs, dim1) return class_names[idx.item()], prob.item()unsqueeze(0)是最容易被忽略的地方。模型forward期望的是4D张量(N, C, H, W)单张图片读取后只有3D尺寸必须补上batch维度。最终softmax在dim1上进行得到每个类别的概率分布后取最大值对应的类这就是置信度最高的预测。源码项目里如果配了gradio或flask接口最终调用的也是这个函数主体。5.2 混淆矩阵绘制与错误样本定位单纯看val_acc会掩盖类别之间的差异。分批次的垃圾类别中可回收物和有害垃圾在塑料材质外观上相近准确率往往最低。用sklearn的confusion_matrix把预测与真实标签画出来第一时间就能看到哪些类别互相混淆。import numpy as np from sklearn.metrics import confusion_matrix, classification_report import seaborn as sns import matplotlib.pyplot as plt model.eval() all_preds [] all_labels [] with torch.no_grad(): for images, labels in val_loader: images images.to(device) outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) print(classification_report(all_labels, all_preds)) # 画图时把类别名标上 sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelstrain_dataset.classes, yticklabelstrain_dataset.classes) plt.show()混淆矩阵的对角线值越高说明每个类别的特征分得越开。查看被错分的样本图时重点是看图片的背景信息——垃圾图像里大量样本的背景是桌面、地面、手部如果某一类样本背景都集中在同一种场景CNN很容易把背景特征当作类别特征。处理办法是增强里加RandomResizedCrop(scale(0.6, 1.0))强制模型关注目标区域而不是背景这也是在原始训练参数基础上的有效进阶调整。5.3 样本量不足时冻结前层的小技巧当某个类别图片只有几十张、训练准确率始终上不去时一个实用技巧是冻结第一个卷积块只优化后面的层。将requires_grad置为False后Adam只更新后半部分参数可学习参数量直接缩减约10%过拟合速度明显变慢。for name, param in model.named_parameters(): if name.startswith(features.0): param.requires_grad False冻结后第一个卷积块仍然参与前向计算但梯度不再回传。这套做法适合小样本迁移场景冻结浅层的边缘、颜色特征提取能力同时释放深层的类别判别能力在垃圾分类这种细粒度不高的任务上收益明显。要快速验证模型加载是否正常用torch.save存权重后执行一次加载比对验证集的准确率是否有差异即可确认推理链路没有断点。本文还有配套的精品资源点击获取
