多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

智能垃圾分类系统:基于PyTorch的图像分类与迁移学习实战

智能垃圾分类系统:基于PyTorch的图像分类与迁移学习实战 简介面向人工智能课程设计、毕业设计及自学实践这份资源提供了一套完整的智能垃圾分类系统源码。系统以图像识别为核心利用卷积神经网络对垃圾图片进行分类覆盖从数据预处理、模型搭建、训练到预测评估的完整流程适合希望掌握AI工程落地要点的学生与开发者。压缩包共52个文件其中28个Python脚本是主体负责数据加载、模型定义、训练与推理等任务另有13个pyc编译文件、6张效果展示图、1份README说明文档及配置JSON等整体体积仅1.15MB目录按功能分化清晰便于定位与二次开发。该资源已有143人学习浏览属于轻量级但可完整运行的实践项目。参考其中的ResNet、CBAM等模型实现可学习如何构建和优化分类网络直接复用训练、预测脚本及配置文件也能为课程作业或毕业设计节省大量搭框架的时间。1. 智能垃圾分类系统毕设和课程作业里最稳的AI落地题智能垃圾分类系统是课程作业和毕业设计里出现频率极高的一类AI视觉项目输入一张垃圾图片模型判断它属于可回收物、厨余垃圾、有害垃圾还是其他垃圾。它受欢迎的原因很实在数据集公开、模型不挑显卡、普通笔记本就能完成训练、最后再套一个简单界面就能做现场演示。标题里“用于人工智能学习所实验”几个字直接点明了这套代码的定位——它是拿来学、拿来改、拿来跑通全流程的不是装在黑匣子里直接交差的。它适合三类人正在选毕设方向的学生、想用一个完整案例走通图像分类流程的入门开发者以及课程项目组里需要在一学期内交付可演示系统的小队。很多人以为难点在模型训练实际做下来你会发现数据组织和演示效果才是真正拉开差距的地方。2. 先想清楚再动手模型选型与数据集准备的常见做法2.1 图片分类还是目标检测毕设场景怎么选做智能垃圾分类第一个决定不是选模型而是选任务形态。如果系统只需要判断“这张照片里的垃圾是什么”图片分类就够用如果场景里有多件垃圾混在一起、需要把每个物体框出来那才需要目标检测。检测的演示效果确实更唬人但标注成本完全不在一个量级——每张训练图片都要人工画边界框公开检测数据集里的垃圾类别和你想做的类别往往对不上自己标又耗时巨大。常见做法是毕设和课程作业默认走分类路线把“智能”体现在类别判断和置信度输出上如果导师明确要求检测再在此基础上换检测框架并且优先找公开的检测数据集而不是自己从零标注。我一般这样判断先问验收场景。如果演示时是“拍一张照片 → 判断这是什么垃圾”分类足够如果是“拍一张桌面 → 指出哪个是可回收物、哪个是有害垃圾”才需要检测。分类模型训练快、改起来直观答辩时也更容易把特征提取、迁移学习这些原理讲透。对多数课程项目来说分类路线已经能把“智能”两个字撑起来没必要一上来就给自己上检测的难度。2.2 数据集的目录组织与标签映射不管用公开数据集还是自己拍先把数据整理成ImageFolder的标准结构train和val两个主目录下面每个类别一个子目录子目录名就是类别名图片直接放在类别目录里。这套结构的好处是torchvision的datasets.ImageFolder可以直接读取不用手写复杂的标签映射逻辑。我这里用一个四分类的例子展示目录怎么组织data/ ├── train/ │ ├── recyclable/ # 可回收物 │ │ ├── plastic_bottle_001.jpg │ │ └── paper_box_002.jpg │ ├── kitchen_waste/ # 厨余垃圾 │ │ ├── banana_peel_001.jpg │ │ └── apple_core_002.jpg │ ├── hazardous/ # 有害垃圾 │ │ ├── battery_001.jpg │ │ └── medicine_002.jpg │ └── other/ # 其他垃圾 │ ├── dirty_rag_001.jpg │ └── ceramic_002.jpg └── val/ ├── recyclable/ ├── kitchen_waste/ ├── hazardous/ └── other/四分类是最常见的设定也有项目把类别细分成六类甚至更多但主体结构不变。有两个细节值得注意子目录名尽量用英文字母不要直接用中文文件夹名否则在Windows上容易触发编码问题后面避坑章节会展开另外类别一旦定下来就不要中途加类每加一个类别所有历史预测结果和评估指标都得重来一遍。类别定好后先跑一段统计脚本看每个类别到底有多少张图。这一步不起眼但能提前暴露类别不均衡的问题。import os from collections import Counter split train # 换成 val 可检查验证集 base fdata/{split} counter Counter() for cls_name in os.listdir(base): cls_dir os.path.join(base, cls_name) if not os.path.isdir(cls_dir): continue n len(os.listdir(cls_dir)) counter[cls_name] n for cls_name, n in counter.most_common(): print(f{cls_name}: {n} 张) min_n min(counter.values()) max_n max(counter.values()) print(f最少/最多: {min_n} / {max_n}, 比值 {max_n / min_n:.2f})这段代码会把每个类别的图片数量打印出来最后一行输出最多类别与最少类别的比值。如果比值超过3训练时模型会对样本多的类别严重偏科后面需要做数据扩充或调整损失函数权重。公开的垃圾分类数据集一般按材料细分比如塑料瓶、玻璃杯、纸盒、果皮、电池这些几十个细类对毕设来说太碎常见做法是把细类归并成四五个大类归并逻辑单独存一份映射文件训练和推理都用同一份映射避免标签错位。2.3 数据增强与类别样本量检查数据增强是这套系统里性价比最高的环节。不需要复杂的生成模型用torchvision自带的transforms就能把有限的数据扩出足够多的多样性。一个典型的训练集增强配置长这样from torchvision import transforms train_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), transforms.ColorJitter( brightness0.2, contrast0.2, saturation0.2, hue0.1 ), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) ]) val_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) ])训练集增强里RandomHorizontalFlip是水平翻转对垃圾这种不区分左右朝向的物体几乎都是安全的RandomRotation给15度以内的旋转用来模拟手机随手拍的倾斜角度ColorJitter调整亮度、对比度、饱和度和色调让模型对不同光线条件下的拍摄结果更鲁棒。验证集不能做随机增强只做缩放和归一化否则同一张图每次验证结果不一样指标失去参考意义。Normalize的mean和std用的是ImageNet预训练统计值只要使用迁移学习这两组数就不要自己改。除了增强我还会顺手做一次样本量检查。若某类只有100张另一个类有900张优先做两件事一是把少样本类别的增强幅度调大比如rotation放宽到30度二是在训练时改用带类别权重的交叉熵损失而不是直接套默认的CrossEntropyLoss。这个权重可以直接按样本数的反比算出来后面训练代码里会体现。3. 训练一个能交差的分类模型完整实现与参数含义3.1 数据加载与预处理从文件夹到DataLoader目录整理好之后加载数据就是标准的ImageFolder加DataLoader两行事。ImageFolder会按子目录名自动生成类别索引比如recyclable对应0、kitchen_waste对应1这个索引顺序就是模型输出的顺序训练和推理必须保持一致。import torch from torch.utils.data import DataLoader from torchvision import datasets batch_size 32 num_workers 2 train_dataset datasets.ImageFolder( rootdata/train, transformtrain_transforms ) val_dataset datasets.ImageFolder( rootdata/val, transformval_transforms ) train_loader DataLoader( train_dataset, batch_sizebatch_size, shuffleTrue, num_workersnum_workers, pin_memoryTrue ) val_loader DataLoader( val_dataset, batch_sizebatch_size, shuffleFalse, num_workersnum_workers, pin_memoryTrue ) print(类别映射:, train_dataset.class_to_idx) print(训练集大小:, len(train_dataset)) print(验证集大小:, len(val_dataset))batch_size32是一个速度和精度比较折中的值显存吃紧就降到16shuffleTrue只在训练集开验证集保持False保证每次评估的数据顺序一致num_workers2在Windows上最稳调大虽然能加快图片读取但在Windows上经常和主进程的调度机制产生兼容性问题表现就是程序卡在第一个batch迟迟不动pin_memoryTrue能稍微减少CPU到GPU的拷贝时间纯CPU环境开着也无妨。注意训练集和验证集最好来自不同的拍摄条件或不同的时间段避免同一张图同时出现在两个集合里。课程项目里最常见的翻车就是同学把收集到的图片直接按比例随机切分验证集和训练集高度相似训练时指标很好看一到现场演示新拍的图片就掉链子。3.2 训练循环优化器、学习率与模型保存模型我一般用ResNet18理由很实际预训练权重好找、参数量小、训练快、对显卡要求低。在公开垃圾分类数据上做迁移学习几轮就能到90%以上的验证准确率。如果机器很老或者计划部署到嵌入式设备可以换成MobileNetV3逻辑完全一样只改模型构造那一行。import os import torch import torch.nn as nn import torch.optim as optim from torchvision import models os.makedirs(checkpoints, exist_okTrue) device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_classes len(train_dataset.classes) model.fc nn.Linear(model.fc.in_features, num_classes) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lr0.001, momentum0.9, weight_decay1e-4) scheduler optim.lr_scheduler.StepLR(optimizer, step_size5, gamma0.1) best_acc 0.0 epochs 20 for epoch in range(epochs): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() scheduler.step() avg_loss running_loss / len(train_loader) 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, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc correct / total print(fepoch{epoch1}, loss{avg_loss:.4f}, val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save({ model_state_dict: model.state_dict(), class_to_idx: train_dataset.class_to_idx, val_acc: val_acc, }, checkpoints/best_model.pth)这段训练代码有几个参数值得细说。优化器用SGD加momentum是迁移学习的常见组合比Adam在这个场景下更容易收敛到平滑的极值点初始学习率0.001是加载预训练权重后的安全值不要一上来就设0.1否则预训练特征会被大幅破坏weight_decay1e-4做L2正则防止小数据集上过拟合。StepLR每5个epoch把学习率降为原来的1/10让收敛后期走得更细。epoch设20是起步值实际训练时盯着验证准确率不再上升就停。模型只在验证准确率超过历史最优时才保存并且把class_to_idx一起存进checkpoint这样推理阶段不需要重新对齐标签。这里有个常见误用如果把验证准确率作为保存模型的依据就不要在训练集上反复调学习率或者无限拉长epoch。见过有人把epoch设到100验证准确率早在第15轮就到平台期后面只是在拟合训练集里的噪声。对课程项目来说训练时间控制在半小时到一小时之间是合理的超过两小时说明数据或超参有地方不对别硬扛。3.3 评估指标只看准确率的毕设容易翻车垃圾分类数据集天然存在类别不均衡可回收物和厨余垃圾的公开图片数量通常远多于有害垃圾。这时候全局准确率会掩盖小类别的失败。举个例子有害垃圾只占5%模型把所有图片都预测成可回收物全局准确率照样能到95%但答辩时一演示有害垃圾就翻车。所以要加两个东西混淆矩阵和每类别的精确率、召回率。from sklearn.metrics import confusion_matrix, classification_report 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.numpy()) class_names list(train_dataset.class_to_idx.keys()) cm confusion_matrix(all_labels, all_preds) print(混淆矩阵行真实类别列预测类别) print(cm) print(classification_report(all_labels, all_preds, target_namesclass_names))这一段把验证集全部数据做一次推理输出语义清晰的混淆矩阵和分类报告。看报告时重点盯每个类别的recall如果某一类明显低于其他类说明模型对这个类别识别能力弱下一步优先给它加数据或调权重而不是盲目换网络结构。classification_report里每一行的recall比全局准确率诚实得多答辩时把这份报告打印出来展示比只报一个“准确率95%”要耐看得多。4. 从模型到演示系统推理封装与可视化界面4.1 单张图片推理脚本去掉训练代码后的最小推理函数训练结束后把推理代码和训练代码分开。训练代码里有反向传播、数据加载、学习率调度这些在现场演示时只会增加出错概率。推理脚本应该短到能在任何一台机器上独立运行加载checkpoint对输入图片做预处理输出类别和置信度。import torch from torchvision import transforms, models from PIL import Image device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet18(weightsNone) checkpoint torch.load(checkpoints/best_model.pth, map_locationdevice) model.fc torch.nn.Linear(model.fc.in_features, len(checkpoint[class_to_idx])) model.load_state_dict(checkpoint[model_state_dict]) model.to(device) model.eval() idx_to_class {v: k for k, v in checkpoint[class_to_idx].items()} infer_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) def predict(image_path: str): img Image.open(image_path).convert(RGB) tensor infer_transforms(img).unsqueeze(0).to(device) with torch.no_grad(): logits model(tensor) probs torch.softmax(logits, dim1).cpu().numpy()[0] pred_idx int(probs.argmax()) pred_class idx_to_class[pred_idx] confidence float(probs[pred_idx]) return pred_class, confidence if __name__ __main__: cls, conf predict(test_images/banana_peel.jpg) print(f预测类别: {cls}, 置信度: {conf:.2%})这里有两个关键点。推理前必须调用model.eval()它会把dropout和BatchNorm切换到推理模式否则同一张图每次预测结果可能不一样——这是新人最容易漏的一行。另一个是加载checkpoint时指定map_locationdevice这样在只有CPU的机器上也能加载之前在GPU上训练的权重。idx_to_class直接从checkpoint里的class_to_idx反转得到而不是在脚本里手写一份。推理脚本里如果自己重新定义类别顺序一旦和训练时的目录排列不一致所有预测结果都会错位。置信度在垃圾分类里非常重要。模型对玻璃瓶和陶瓷碎片这种外观接近的物体经常给出60%左右的置信度——这不是模型坏了是类别边界本身就模糊。演示时把置信度直接显示在界面上能有效降低“模型答错”的观感因为使用者能看到模型的不确定性。4.2 用Web界面把模型变成可演示系统课程作业的验收场景通常是现场演示命令行打印结果虽然专业但不够直观。常见做法是套一个Web界面上传图片点击识别页面显示垃圾类别、置信度以及对应的投放建议。我用Streamlit做这套界面原因是代码量少、一个文件能写完、不需要任何前端经验答辩前临时改布局也不慌。import streamlit as st from PIL import Image st.set_page_config(page_title智能垃圾分类系统, layoutcentered) st.title(智能垃圾分类系统) uploaded st.file_uploader(上传垃圾图片, type[jpg, jpeg, png]) if uploaded is not None: img Image.open(uploaded).convert(RGB) st.image(img, caption待识别图片, use_column_widthTrue) cls, conf predict(uploaded) st.success(f识别结果{cls}置信度 {conf:.1%}) advice { recyclable: 投放进可回收物桶注意清空残留液体。, kitchen_waste: 投放进厨余垃圾桶建议沥干水分。, hazardous: 投放进有害垃圾桶避免破损。, other: 投放进其他垃圾桶。, } st.info(advice[cls])这段界面直接把predict函数从推理脚本里import进来即可逻辑就是上传图片、显示图片、调用预测、渲染结果和投放建议。use_column_widthTrue让图片自适应列宽避免大图把页面撑变形。st.success和st.info是Streamlit的状态组件用颜色区分结果和提示演示时视觉效果清楚。摄像头场景可以再加一个OpenCV取帧的循环把每一帧传给predict。这里有一个性能要点推理一次大约几十到几百毫秒取决于机器摄像头场景要控制抽帧频率一般每秒处理2到3帧就够了而不是每帧都推理——后者会把界面卡死看起来系统极不稳定。抽帧的常见做法是记录上一次推理的时间戳超过0.3秒才处理新帧或者用队列丢帧保证界面始终能响应。注意如果现场网络不稳定Streamlit首次运行会自动检查更新可能卡在启动界面。答辩前先离线跑一次或者在启动命令里加环境变量禁用自动检查别让这个细节毁掉演示。5. 智能垃圾分类系统常见问题排查与避坑记录5.1 训练loss正常下降但验证准确率不动现象训练集loss从2.0稳步降到0.2验证准确率却一直卡在某个值附近不动比如稳定在45%。原因这多半不是模型问题是数据集的目录结构或类别顺序出了问题。常见的是验证集和训练集的类别分布不一致或者ImageFolder在两个集合里生成的类别索引不同——train里recyclable是0val里recyclable是1模型在训练时学的0号类别对应验证集的1号类别准确率自然上不去。解决在训练脚本和评估脚本里分别打印train_dataset.class_to_idx和val_dataset.class_to_idx确认两个映射完全一致。更彻底的办法是只创建一个datasets.ImageFolder(rootdata)实例然后按文件路径手动切分训练集和验证集保证类别索引只生成一次两边天然对齐。5.2 Windows路径与中文标签导致的数据加载失败现象代码在Linux上跑得好好的换到Windows上训练直接报错提示找不到图片文件或者ImageFolder读出来的类别名变成乱码。原因Windows默认编码是GBK而公开数据集的类别标签大多是UTF-8编码中文类别名在部分PyTorch版本上还会触发编码异常。路径分隔符不同也会让硬编码的路径失效比如在Linux上写死的“data/train”在Windows上可能没问题但带中文的绝对路径就很容易翻车。解决所有类别目录统一用英文字母命名代码里的路径用os.path.join拼接不手写“/”或“\”读任何CSV或JSON标签文件时都显式指定encodingutf-8。这些改动很小但能省掉换机器时的半天排查时间。5.3 显存不足与batch size的连锁问题现象训练到中途报CUDA out of memory程序直接崩溃。把batch_size从32调到8后能跑但训练变慢而且loss震荡明显变大。原因显存溢出是输入尺寸或batch太大但直接把batch_size调小会改变BatchNorm的统计行为小batch下每个batch的均值和方差噪声变大训练稳定性下降。解决先保持batch_size不变把输入resize从224下调到160显存立刻省出一大截分类任务对输入尺寸没那么敏感如果必须用batch_size8就把学习率从0.001同步降到0.0003左右并且考虑用梯度累积——每4个step累加一次梯度再更新效果接近batch_size32。梯度累积的写法不复杂用一个变量累积loss.backward()后的梯度每N步才optimizer.step()一次。5.4 摄像头推理卡顿与画面延迟现象摄像头画面只有每秒几帧点击识别后界面完全卡住过好几秒才出结果。原因每一帧都做了完整推理CPU推理一张224x224的图片通常需要100到300毫秒再加上画面渲染交互自然卡死。解决抽帧推理每秒最多处理2到3帧其他帧直接丢弃更稳健的做法是把推理放到独立线程界面主线程只负责显示最近一次的推理结果。演示时还要提前准备几张测试图放在桌面避免现场打开网页找图片浪费时间。5.5 换数据集后标签映射错位现象之前验证准确率不错的模型换了一批新图片测试结果全部分到同一个类别或者类别名对不上。原因保存模型时只存了模型参数没存类别映射。推理脚本里手写了一份class_to_idx但顺序和新数据集的文件夹顺序不一致导致所有输出都被错误映射。解决训练时把class_to_idx一起存进checkpoint推理时直接从checkpoint读取不在推理脚本里二次定义。前面训练代码里已经做了这一步但很多人图省事会在推理脚本里手写一份映射就是给自己埋坑。另外换数据集后第一件事是重新打印checkpoint里的class_to_idx确认类别顺序没变再跑推理。6. 让答辩从“能跑”变成“能讲”两个百试百灵的演示技巧6.1 混淆矩阵与分类报告比一行准确率更能说明问题答辩时评委一定会问“你的系统效果到底怎么样”。只回答“准确率95%”这个数字孤立无援但把混淆矩阵和classification_report放进演示文档就能指着每一行讲哪个类别识别得好、哪个类别容易混淆、为什么混淆。讲出“玻璃和陶瓷外观相似所以互相误判较多”这句话比重复十遍“我用了ResNet18”有说服力得多。做法也简单提前把验证集的分析结果截图存好另外挑3到5张典型误判图片分析模型为什么错。一个很容易踩的坑是误判图片从网上下载和训练集分布差异太大现场展示反而暴露问题。正确做法是拿自己验证集里的失败case讲这样能体现你对数据分布的理解。如果某个类别召回率特别低还可以现场给出对策——给该类别补充样本、增加增强强度、或在损失函数里加类别权重这些“下一步改进方向”就是答辩加分项。6.2 用类激活图证明模型真的在“看”垃圾第二个实用技巧是可视化模型的关注区域思路来自Grad-CAM取模型最后一个卷积层的梯度加权得到一张和输入图片同尺寸的热力图标出模型分类时主要在看图片的哪个部位。垃圾分类场景下一个合格的模型应该关注垃圾主体本身而不是背景桌面或手掌。如果热力图显示模型盯着图片角落的阴影说明模型学到了数据集里的背景噪声也就是常说的“模型魔怔了”——这种发现拿到答辩现场比单纯展示“准确率98%”更能体现你对模型行为的理解。这个可视化脚本不复杂原理上就是一次反向传播取梯度再和特征图逐通道加权求和。它不需要改模型结构适配现有训练好的checkpoint就能跑是典型的低成本高产出功能。我一般会把几张典型图片的热力图和原图并排保存成一张长图演示时一页就能讲清楚“模型在看什么”。我自己的习惯是答辩前一周把模型、界面、分析脚本全部冻结不再改任何超参只反复走三遍完整演示流程——上传图片、看结果、点开置信度。垃圾分类这个题目做到“能讲清楚原理、能指出失败案例、能给出下一步改进方向”就已经超过大多数只报准确率的课程作业了。希望帮到你。本文还有配套的精品资源点击获取
返回列表