ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

PyTorch垃圾分类实战:从CNN到ResNet迁移学习与精度调优

2026/10/2 9:58:51 拓冰建站 浏览量
PyTorch垃圾分类实战:从CNN到ResNet迁移学习与精度调优 简介这套垃圾分类实战项目以卷积神经网络CNN和ResNet残差网络为核心面向深度学习初学者和计算机视觉爱好者提供从模型搭建、训练到评估的完整流程。压缩包内共5个文件包括两个Python源码脚本、两个预训练模型权重文件和一个验证集压缩包整体仅4MB轻量易部署。目前已有273人学习下载。代码覆盖数据预处理、数据增强、模型构建、交叉熵损失训练、学习率与Dropout调参、准确率和混淆矩阵评估等关键环节并附有可直接运行的预测脚本。训练好的权重文件可在PyTorch环境中加载便于直接测试垃圾分类效果通过阅读源码还能深入理解CNN特征提取、ResNet残差连接缓解深层网络退化的原理。整体目录结构按源码、权重和数据集分层组织方便按需查阅适合作为课程设计或毕业设计的参考模板。1. 垃圾分类实战为什么CNN不够用要上ResNet先说结论用纯CNN做垃圾分类跑通流程没问题但换到真实场景里你会发现它在纹理相近的类别上经常翻车——比如玻璃瓶和陶瓷碗普通卷积网络学到的特征不够“细”。这套基于CNN和ResNet的垃圾分类实战资源解决的就是从“模型能跑”到“模型能判别”这一段路。它适合两类人一是课程作业或毕设选了这个方向、需要一个完整可复现流程的在校生二是刚接触深度学习图像分类、想搞明白卷积网络和残差网络到底差在哪的Python工程师。资源里的代码以PyTorch为主从数据处理、模型构建到训练调参都有对应脚本你不需要懂太多原理就能把结果跑出来但每一步的选型理由我会在下面逐个拆开讲。2. 数据准备与预处理从图片目录到可训练Tensor2.1 数据集结构与类别映射垃圾分类公开数据集最常见的形态是一个根目录下面按类别分文件夹每个文件夹里是同类的图片。以华为云的四十类垃圾分类数据集和Kaggle上的垃圾分类数据集为例前者图片数量大、类别细后者轻量、适合跑通流程。这套实战项目默认你拿到的是这种目录结构garbage_data/ ├── battery/ # 电池有害垃圾 ├── clothes/ # 衣物可回收 ├── glass_bottle/ # 玻璃瓶可回收 ├── kitchen_waste/ # 厨余垃圾 └── ... # 其余类别如果你的数据集不是这种结构或者图片文件名特别乱我一般会先写一个小脚本统一成上面这种格式。因为PyTorch的torchvision.datasets.ImageFolder天然依赖这种按类别分文件夹的组织方式它能根据文件夹名自动生成类别索引省掉手写标签映射表的功夫。2.2 数据增强与归一化参数数据增强这块直接照搬ImageNet那套默认参数容易踩坑——后面准确率上不去很有可能就是预处理和模型不匹配。CNN基线模型和ResNet预训练权重的输入协议不同ResNet加载的是ImageNet预训练权重它期望输入图片被归一化到ImageNet数据集的均值方差分布而自己从头训练的CNN则没那么挑。这套实战里对两者的预处理做了区分# 训练集增强先随机裁剪再缩放模拟不同拍摄距离 train_transforms transforms.Compose([ transforms.RandomResizedCrop(size224, scale(0.8, 1.0)), 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], # ImageNet均值 std[0.229, 0.224, 0.225]) # ImageNet方差 ]) # 验证集/测试集不增强只做缩放和归一化 val_transforms transforms.Compose([ transforms.Resize(size256), transforms.CenterCrop(size224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])参数说明RandomResizedCrop的scale(0.8, 1.0)控制裁剪面积占原图的比例一般不建议低于0.5否则会裁掉太多关键信息RandomRotation(15)只旋转15度因为垃圾分类图片不像遥感图像那样需要大角度旋转。ColorJitter三个参数分别控制亮度、对比度和饱和度值设得太大容易让模型学到颜色噪声。ResNet迁移学习时归一化参数必须用ImageNet的这套固定值偏差太大会让预训练权重产生“不适应”。2.3 自定义Dataset与DataLoader虽然ImageFolder可以直接用但实际项目中我更推荐自己写一个Dataset类。原因很实际垃圾图片里有不少是非标准尺寸的、带EXIF旋转信息的手机照片ImageFolder遇到这类图会直接报错而自定义Dataset可以顺手做容错处理。class GarbageDataset(Dataset): def __init__(self, root_dir, transformNone): self.classes sorted([d for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d))]) self.class_to_idx {cls: idx for idx, cls in enumerate(self.classes)} self.images [] for cls in self.classes: cls_dir os.path.join(root_dir, cls) for img_name in os.listdir(cls_dir): if img_name.lower().endswith((.jpg, .jpeg, .png)): self.images.append((os.path.join(cls_dir, img_name), self.class_to_idx[cls])) self.transform transform def __len__(self): return len(self.images) def __getitem__(self, idx): img_path, label self.images[idx] try: image Image.open(img_path).convert(RGB) # 统一转成RGB三通道 except Exception: # 遇到损坏图片返回同类别的一张替代 return self.__getitem__((idx 1) % len(self.images)) if self.transform: image self.transform(image) return image, label这个类的核心设计有三个第一加载时用一个列表把路径和标签固化下来避免每次迭代都做磁盘扫描第二convert(RGB)强制统一通道解决灰度图和RGBA图混入的问题第三__getitem__里做异常兜底损坏图片直接跳到下一张训练过程不会因为单张坏图中断。DataLoader参数也不能全用默认值。batch_size在单卡GPU上建议32起步显存不够就降到16num_workers在Windows上设成0或者2设大了反而容易报错pin_memoryTrue能减少数据从CPU拷贝到GPU的开销训练速度有可感知的提升。3. 模型搭建CNN基线模型与ResNet迁移学习双轨实现3.1 CNN基线模型三层卷积足以跑通全流程这套资源里先给了一个轻量CNN作为基线目的不是拿它打榜而是让你先确认数据链路、训练脚本、评估逻辑都是通的。如果一开始就上ResNet出了问题很难分清是模型问题还是数据问题。基线的结构很简单三层卷积加两个全连接层。import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes40): super(SimpleCNN, self).__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), # 输出 32x224x224 nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), # 输出 32x112x112 nn.Conv2d(32, 64, kernel_size3, padding1), # 输出 64x112x112 nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), # 输出 64x56x56 nn.Conv2d(64, 128, kernel_size3, padding1), # 输出 128x56x56 nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2) # 输出 128x28x28 ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(128 * 28 * 28, 256), nn.ReLU(inplaceTrue), nn.Dropout(p0.5), nn.Linear(256, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x逻辑说明每个卷积层后面跟一个BatchNorm再激活这是比单纯ConvReLU更稳的组合能避免深层网络梯度消失。Dropout(0.5)放在第一个全连接层后专门用来压过拟合——垃圾图片的纹理、颜色、背景分布很复杂模型很容易记住训练集里的环境噪声而不是垃圾本身。参数说明kernel_size3, padding1的组合保证卷积不改变特征图尺寸尺寸变化完全由MaxPool2d控制。输入224x224的图经过三次下采样后变成28x28这个尺寸对全连接层来说计算量还能接受。如果你要识别的类别数不是40改num_classes参数即可。这个基线的训练集准确率通常能到92%以上但验证集一般在85%左右徘徊——这就是纯CNN的特征表达力瓶颈。不是代码有bug是模型容量不够这时再增加卷积层数边际收益也很低应该切换到ResNet。3.2 ResNet迁移学习加载预训练权重并替换分类头迁移学习的核心思路是ImageNet预训练模型已经把通用视觉特征边缘、纹理、形状学好了我们只需要替换最后的分类层让它在垃圾分类这个特定任务上重新学习。这套实战里用torchvision.models.resnet18做主力显存充足的情况下可以换成resnet50精度会更高但训练时间大约翻三倍。import torchvision.models as models def build_resnet(num_classes40, pretrainedTrue): model models.resnet18(pretrainedpretrained) # 冻结前四层卷积参数只训练最后几层 for name, param in model.named_parameters(): if layer4 not in name and fc not in name: param.requires_grad False # 替换最后一层全连接输出维度改为类别数 in_features model.fc.in_features model.fc nn.Sequential( nn.Dropout(p0.3), nn.Linear(in_features, 256), nn.ReLU(inplaceTrue), nn.Linear(256, num_classes) ) return model参数说明pretrainedTrue会从网上下载ImageNet预训练权重第一次运行需要保持网络通畅之后会缓存到本地。requires_gradFalse表示冻结参数被冻结的层只做前向计算不参与反向传播这样训练参数量大幅减少显存占用也会降低。冻结策略这里有个经验如果数据集和ImageNet图像差异大比如医学影像、遥感图冻结太狠效果反而差但垃圾分类图片和ImageNet里的日常物品重合度比较高冻结浅层是安全的。这套实战里把layer4之前的全部冻结layer4和新增的分类头参与训练是一个比较折中的方案。3.3 损失函数与优化器选择分类任务的标准配置是交叉熵损失函数加Adam优化器这套实战里也没有搞特殊化。但有一个细节很重要CrossEntropyLoss在PyTorch里已经内置了Softmax操作所以模型最后一层不需要额外加Softmax激活。criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr0.001, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max30, eta_min1e-5 )逻辑说明weight_decay1e-4是L2正则化用来惩罚过大权重对防止过拟合有实际帮助尤其在数据量不足一万的情况下。学习率调度选用了余弦退火而不是固定学习率因为余弦退火在训练后期能以更小的学习率微调权重最后几轮往往能带来一两个百分点的提升。训练循环里还需要做一个关键动作每个epoch结束后在验证集上计算准确率只在准确率提升时保存模型权重。这个叫“基于验证集的最优模型保存”比盲目保存最后一个epoch的模型靠谱得多。资源里的训练脚本已经写好了这个逻辑你只需要跑起来就能看到训练曲线和验证曲线。提示如果验证集准确率在某个阈值附近不断震荡不用急着改模型结构先检查学习率是不是太大尝试把初始学习率降到0.0003。4. 训练避坑与常见问题六个典型翻车现场4.1 图片读不进来全是灰的或者报“Found 0 images”现象训练脚本一启动就提示找到0张图片或者loss一直是NaN图表里全是空的。原因数据集里混入了非图片文件或者文件夹嵌套层级不对ImageFolder找错了路径。解决在加载数据前用脚本做一次文件体检把所有文件的扩展名和尺寸都打印出来确认最深的路径是否到达图片文件。我一般直接在Dataset初始化时加一行调试代码assert len(self.images) 0, fNo images found in {root_dir}如果运行到这里就中断说明路径结构有问题这时先修数据目录而不是改模型。4.2 训练时报错“Expected 3D tensor but got 2D”现象数据加载时报Expected 3D tensor but got 2D或者模型forward时报维度不匹配。原因数据集里混有灰度图灰度图只有两个维度H, W而模型期望的是三通道C, H, W。这种情况在垃圾分类数据集里特别常见因为有些旧手机拍出的照片或扫描件确实是灰度模式。解决在第2章的自定义Dataset里强制convert(RGB)或者在进入模型前统一做img img.repeat(3, 1, 1)把单通道复制成三通道。4.3 加载预训练权重失败URL下载超时或证书报错现象pretrainedTrue时下载卡在99%不动或者报SSL证书错误。原因国内网络访问国外模型权重库不稳定这是环境问题不是代码问题。解决手动下载权重文件到本地然后通过load_state_dict加载model models.resnet18(pretrainedFalse) state_dict torch.load(resnet18-f37072fd.pth, map_locationcpu) model.load_state_dict(state_dict)权重文件用搜索引擎搜“resnet18权重下载”能找第三方镜像注意比对文件MD5值防止下到损坏文件。4.4 准确率上不去基线比ResNet高迁移学习反而更差现象从头训练的CNN验证集准确率有88%换ResNet后反而只有80%。原因这是典型的“预训练权重被破坏”问题。最常见的原因是修改全连接层之后再加载权重或者训练时冻结策略选错了层——把不合适的层也冻结了导致ResNet只剩分类头在训练而分类头的随机初始化参数又要匹配前面固定特征。解决先确认训练脚本里是“先加载权重再替换结构”如果本地有多个权重文件检查文件是否完整。还可以试试把pretrainedFalse跑一次对比如果准确率一样说明预训练权重根本没有加载进去或已损坏。4.5 显存溢出和训练速度缓慢现象CUDA out of memory报错或者一个epoch要跑十几分钟。原因图片分辨率是224x224batch_size设到64时要占约4GB显存带特征冻结的ResNet还好换成从头训练的ResNet就直接爆了。解决把batch_size降到16或8同时把num_workers调低Windows系统建议设为0。如果还想提速把输入尺寸从224降到160试试训练速度能提升接近一倍精度损失通常在1到2个百分点以内。4.6 数据类别不平衡某些垃圾种类永远预测不对现象训练完看每类准确率电池、灯管这类图片数量少的类别准确率不到30%总量大的类别却有98%。原因垃圾分类数据集天然不平衡厨余垃圾图片数量远远多于有害垃圾。解决先统计每类图片数量再做两类操作——对少数类做过采样重复读图或对多数类做欠采样同时给损失函数按类别数量设置权重weights torch.tensor([1.0 / count for count in class_counts]) criterion nn.CrossEntropyLoss(weightweights.to(device))这段代码里class_counts是一个列表顺序要和classes文件夹顺序保持一致。加权重后模型会提高对少数类的惩罚力度预测倾向会更均衡。5. 精度调优与部署验证把准确率从91%提到96%5.1 分层冻结微调策略当验证集准确率进入平台期常规训练已经不起作用时我一般会做“解冻重训”把之前冻结的层全部解冻但给它们一个很小的学习率比如1e-5只让参数做微调不破坏已学到的特征。这套实战里可以用参数分组来实现不同模块对应不同学习率optimizer torch.optim.Adam([ {params: model.layer1.parameters(), lr: 1e-5}, {params: model.layer2.parameters(), lr: 1e-5}, {params: model.layer3.parameters(), lr: 1e-4}, {params: model.layer4.parameters(), lr: 1e-4}, {params: model.fc.parameters(), lr: 1e-3}, ], weight_decay1e-4)这个设计的逻辑是浅层特征通用性强学习率要小深层特征与任务关联度高学习率要大。最忌讳的做法是全部层用同一个学习率重新训练那样会把预训练权重破坏得干干净净。5.2 类别不平衡时改用Focal Loss如果加权重交叉熵依然表现不佳Focal Loss是更进一步的方案。它专门针对难分类样本设计通过调整聚焦参数让模型更关注“分错”的样本。对于垃圾分类场景那些包装被压扁、颜色发生变化的物品就是天然难样本class FocalLoss(nn.Module): def __init__(self, gamma2.0, alpha0.25): super(FocalLoss, self).__init__() self.gamma gamma self.alpha alpha def forward(self, inputs, targets): ce_loss nn.functional.cross_entropy(inputs, targets, reductionnone) pt torch.exp(-ce_loss) focal_loss self.alpha * (1 - pt) ** self.gamma * ce_loss return focal_loss.mean()参数说明gamma2.0控制对易分样本的降权力度gamma越大越关注难样本但设太大训练容易不稳定alpha0.25是正负样本平衡系数在二分类里常用多分类时可以按类别出现频率倒数重设。5.3 用混淆矩阵定位“长得像”的类别调优到了后半段光看总体准确率已经不够了我会打印混淆矩阵看看到底哪些类互相混淆。实践中一个高频规律是厨余垃圾里的“剩饭”和“菜叶”互相混可回收垃圾里的“纸盒”和“报纸”互相混。原因不是模型笨而是这两类物体在颜色、形状和拍摄角度上确实接近。矩阵可视化不一定要写很复杂的代码两行就能出图from sklearn.metrics import confusion_matrix, classification_report import matplotlib.pyplot as plt cm confusion_matrix(all_labels, all_preds)打印classification_report看每一类的Precision和Recall低于80%的类别单独提出来看样本图。你会发现很多错分样本里物体的关键特征被遮挡或背景过于杂乱。处理方式有两种一是增加这类图片的数量从原始数据里把对应类别的所有图片都加入训练集扩充二是针对这类物体单独做一个二分类微调而不是反复折腾整个模型。最后说说我一个踩了两次才记住的教训每次跑迁移学习都先用一小批数据比如每类20张图快速跑5个epoch验证数据链路是否通畅确认没问题再启动全量训练。这个“小跑验证”的步骤能帮你省下大量等待时间我第一次直接全量训练跑了两个小时才发现数据没做归一化等于白跑。从那以后凡是新数据集第一次喂给新模型我都会强制走一遍小批量验证加打印loss曲线的流程。希望帮到你这套实战里该省的时间点我已经替你提前避开了。本文还有配套的精品资源点击获取