ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

基于迁移学习的Python垃圾分类系统:从课程设计到端到端落地

基于迁移学习的Python垃圾分类系统:从课程设计到端到端落地 简介这份资源是面向高校学生与机器学习初学者的课程设计级项目源码围绕Python垃圾分类系统展开帮助读者在真实场景中理解监督学习从数据到部署的完整链路。压缩包共32个文件约2.26MB以jpg、jpeg、png图像样本与xml标注为主辅以4个py脚本、xls评估结果及iml工程配置覆盖数据、训练与结果记录各环节。项目涉及scikit-learn、numpy、pandas、OpenCV等库内容涵盖图像预处理、颜色与纹理特征工程、PCA降维、决策树/随机森林/SVM/KNN及MobileNet迁移学习等模型训练并通过准确率、召回率、F1分数与混淆矩阵完成评估最终借助Web框架实现图片上传识别。已有772人学习适合作为课程设计参考或机器学习入门实战模板便于快速复现与二次开发。1. 从一份「机器学习课程设计」说起Python 垃圾分类系统到底在做什么期末周前两周实验室里最常见的一幕是有人抱着「机器学习课程设计」的题目发愁选题写着「Python 垃圾分类系统源码.zip」打开一看要么跑不起来要么准确率惨不忍睹要么干脆只有一堆.py文件没有权重。这个标题背后其实是一个很典型的落地任务——用 Python 搭一套能识别垃圾类别的图像分类系统输入一张照片输出「可回收物 / 厨余垃圾 / 有害垃圾 / 其他垃圾」四类中的一类。它解决的不是「垃圾分类」这个社会问题本身而是让你在有限时间内跑通一条完整的机器学习流水线数据准备、模型训练、推理封装、界面展示。适合谁适合正在做机器学习课程设计、想拿一个能演示能答辩的项目、又不想从零造轮子的同学也适合刚入门想找一个端到端小项目练手的开发者。下面我按自己踩过的路把这条流水线拆开讲清楚。2. 先想清楚技术选型为什么是迁移学习而不是从零训 CNN2.1 四分类任务的真实难度与数据规模垃圾分类图像分类看着简单实际有几个坑类间差异小比如用过的纸巾到底算厨余还是其他各地标准还不一样类内差异大一个塑料瓶可以是完整的、压扁的、带标签的背景杂乱拍摄环境从桌面到垃圾桶都有。如果你打算从零训练一个 CNN至少需要每类几千张标注图还要调很久。课程设计的时间通常只有一到两周自己爬数据、清洗、标注根本来不及。所以主流做法是迁移学习拿一个在 ImageNet 上预训练好的骨干网络冻结大部分层只训练最后的分类头。常见选择是 ResNet18/50、MobileNetV3、EfficientNet-B0。ResNet18 参数量约 11M单张 1080Ti 上推理几毫秒课程设计够用MobileNetV3 更轻适合后面要部署到树莓派或手机的场景。2.2 数据集从哪来公开集 自采的混合策略公开的垃圾分类数据集质量参差常见的有 TrashNet6 类约 2500 张、华为云垃圾分类大赛数据、以及一些 Kaggle 上的中文垃圾数据集。我的建议是先用公开集把流程跑通再补拍 200300 张自己场景的图做微调。目录结构按ImageFolder的约定来这是最省事的dataset/ ├── train/ │ ├── recyclable/ # 可回收物 │ ├── kitchen/ # 厨余垃圾 │ ├── hazardous/ # 有害垃圾 │ └── other/ # 其他垃圾 └── val/ ├── recyclable/ ├── kitchen/ ├── hazardous/ └── other/每个类别文件夹名就是标签名torchvision.datasets.ImageFolder会自动按文件夹名映射成 0/1/2/3。注意类别顺序是按文件夹名字母序排的hazardous会排在kitchen前面后面写推理代码时标签列表必须和这个顺序一致否则会出现「模型说可回收、界面显示有害」的玄学 bug。2.3 训练脚本冻结骨干 替换分类头的最小实现下面这段是能直接跑的核心训练代码基于 PyTorch。关键点是加载预训练权重、替换fc层、只优化分类头参数。import torch import torch.nn as nn from torchvision import datasets, models, transforms from torch.utils.data import DataLoader # 数据增强训练集用随机裁剪翻转验证集只做 resize 和归一化 train_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) val_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) train_ds datasets.ImageFolder(dataset/train, transformtrain_tf) val_ds datasets.ImageFolder(dataset/val, transformval_tf) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4) # 加载预训练 ResNet18替换最后的全连接层为 4 分类 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) for param in model.parameters(): param.requires_grad False # 先冻结全部 model.fc nn.Linear(model.fc.in_features, 4) # 新层默认 requires_gradTrue device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.fc.parameters(), lr1e-3) for epoch in range(15): model.train() for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() loss criterion(model(imgs), labels) loss.backward() optimizer.step() # 每轮验证 model.eval() correct total 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) pred model(imgs).argmax(1) correct (pred labels).sum().item() total labels.size(0) print(fepoch {epoch}, val_acc{correct/total:.4f}) torch.save(model.state_dict(), trash_resnet18.pth)逻辑说明requires_gradFalse冻结骨干只有新fc层参与训练这样 15 轮就能收敛显存占用也小。lr1e-3是分类头常用的起点如果验证准确率震荡就降到 3e-4。batch_size32在 8G 显存上跑 224×224 没问题显存不够就降到 16。归一化用的均值方差是 ImageNet 的标准值别自己乱改否则预训练权重的特征分布对不上准确率会掉一截。训练完保存的是state_dict推理时要先重建同样的网络结构再load_state_dict。3. 推理封装与界面让答辩现场能点一下就跑3.1 单张图片推理函数与标签顺序对齐训练完只是第一步答辩时老师要看的是「我传一张图它告诉我是什么垃圾」。推理脚本要处理好三件事模型结构重建、权重加载、标签顺序。下面这个函数可以直接复用import torch from torchvision import models, transforms from PIL import Image # 标签顺序必须和 ImageFolder 的字母序一致 CLASSES [hazardous, kitchen, other, recyclable] CN_NAME { hazardous: 有害垃圾, kitchen: 厨余垃圾, other: 其他垃圾, recyclable: 可回收物, } def load_model(weight_pathtrash_resnet18.pth): model models.resnet18(weightsNone) model.fc torch.nn.Linear(model.fc.in_features, 4) model.load_state_dict(torch.load(weight_path, map_locationcpu)) model.eval() return model def predict(img_path, model): tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) img Image.open(img_path).convert(RGB) x tf(img).unsqueeze(0) # 加 batch 维度 with torch.no_grad(): prob torch.softmax(model(x), dim1)[0] idx prob.argmax().item() return CN_NAME[CLASSES[idx]], prob[idx].item() if __name__ __main__: m load_model() label, score predict(test.jpg, m) print(f预测{label}置信度 {score:.3f})CLASSES的顺序是血泪经验如果你训练时文件夹叫recyclable/kitchen/hazardous/other字母序排出来是hazardous, kitchen, other, recyclable推理时标签列表写错模型明明预测对了显示出来却是错的这种 bug 能查一下午。map_locationcpu是为了在没有 GPU 的答辩电脑上也能加载。unsqueeze(0)补上 batch 维度因为模型 forward 期望输入是[N, C, H, W]。3.2 用 Gradio 三行代码搭一个能演示的界面课程设计答辩最忌讳只给命令行。用 Gradio 可以几分钟做出一个网页界面老师扫码就能玩import gradio as gr model load_model() def classify(img): # Gradio 传进来的是 numpy 数组转成 PIL from PIL import Image pil Image.fromarray(img) tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) x tf(pil).unsqueeze(0) with torch.no_grad(): prob torch.softmax(model(x), dim1)[0] idx prob.argmax().item() return {CN_NAME[CLASSES[i]]: float(prob[i]) for i in range(4)} gr.Interface(fnclassify, inputsgr.Image(), outputsgr.Label(num_top_classes4), title垃圾分类识别).launch()gr.Label输出的是各类别概率字典界面上会显示成条形图比单纯打印一个类别更有说服力。launch()默认起在127.0.0.1:7860答辩时本机打开浏览器即可。如果要在局域网内让老师用手机访问加server_name0.0.0.0但注意别在公网环境随便开。3.3 准确率上不去的三个调参方向如果验证准确率卡在 70% 左右上不去按这个顺序排查第一检查数据本身用ImageFolder打印每类样本数如果某类只有几十张先补数据或做类别加权第二解冻骨干的后几层一起微调把layer4的requires_grad设为True学习率降到 1e-4通常能涨 510 个点第三换更强的数据增强比如RandomResizedCrop、ColorJitter模拟拍摄时的光照和角度变化。别一上来就换模型ResNet18 在几千张图上足够换大模型反而容易过拟合。4. 避坑与排查那些让课程设计翻车的细节4.1 现象训练 loss 一直不降准确率等于随机猜原因通常是标签和输入没对上或者归一化参数写错。ImageFolder的标签是按文件夹名字母序生成的如果你手动写了个label_map又和它不一致模型学到的就是错位的映射。解决训练前打印train_ds.classes和train_ds.class_to_idx确认顺序推理时直接用这个class_to_idx反查别自己手写列表。4.2 现象验证准确率很高但自己拍的照片全错这是典型的域偏移。公开数据集多是白底、单物体、正面拍摄你自己拍的是桌面杂乱背景、侧面、有阴影。原因在于模型学到了背景特征而不是物体特征。解决补拍 200 张自己场景的图加入训练集或者用RandomResizedCrop加ColorJitter做更强的增强逼模型关注物体本身。另一个技巧是推理时先做一次简单的中心裁剪去掉边缘干扰。4.3 现象CUDA out of memory或者训练到一半卡死batch_size太大或者num_workers设太高。8G 显存跑 ResNet18 224×224batch_size32是安全线batch_size64就可能爆。num_workers在 Windows 上设大于 0 有时会卡在第一个 epoch设成 0 最稳Linux 上设 4 没问题。解决先降batch_size到 16num_workers设 0跑通再往上加。4.4 现象保存的模型加载时报Missing key(s) in state_dict保存和加载的网络结构不一致。常见于训练时用了models.resnet18(weights...)推理时用了models.resnet18(weightsNone)但忘了替换fc层或者反过来多包了一层nn.DataParallel。解决保存时存model.state_dict()而不是整个model加载时先构建完全相同的结构再load_state_dict。如果训练用了DataParallel保存的 key 会带module.前缀加载时用{k.replace(module., ): v for k, v in sd.items()}去掉。4.5 现象Gradio 界面传图后报Expected 4D inputGradio 的gr.Image()默认传进来是 numpy 数组形状是[H, W, C]而模型要的是[N, C, H, W]。解决先Image.fromarray转 PIL再走transforms最后unsqueeze(0)。别直接把 numpy 数组喂给transforms.ToTensor()它虽然也能处理但通道顺序容易出问题。5. 进阶技巧把课程设计做成能写进简历的项目课程设计如果只停在「跑通」答辩完就忘了。想让它变成简历上能讲的东西有两个方向可以加。第一个是模型量化与部署用torch.quantization.quantize_dynamic把模型转成 int8体积缩小到原来的四分之一推理速度提升两三倍然后写一个 Flask 接口用requests压测一下 QPS。这个过程中你会真正理解「训练」和「推理」是两回事。第二个是错误分析把验证集里预测错的图挑出来按类别统计混淆矩阵看看模型到底把「厨余」错认成「其他」还是「可回收」。我一般会写一段脚本把错分图存到errors/目录按真实标签和预测标签命名翻一遍就能发现数据里的标注噪声——有些图我自己都分不清模型错了不冤。# 混淆矩阵 错分图导出 from sklearn.metrics import confusion_matrix import numpy as np, os from PIL import Image model.eval() y_true, y_pred [], [] os.makedirs(errors, exist_okTrue) with torch.no_grad(): for imgs, labels in val_loader: preds model(imgs.to(device)).argmax(1).cpu() for i in range(len(labels)): y_true.append(labels[i].item()) y_pred.append(preds[i].item()) if labels[i] ! preds[i]: # 保存错分图文件名带真实和预测标签 img imgs[i].permute(1, 2, 0).numpy() img (img * [0.229, 0.224, 0.225] [0.485, 0.456, 0.406]) * 255 Image.fromarray(img.astype(uint8)).save( ferrors/true{labels[i]}_pred{preds[i]}_{i}.jpg) print(confusion_matrix(y_true, y_pred))这段代码跑完你会得到一张 4×4 的混淆矩阵和一堆错分图。混淆矩阵告诉你哪两类最容易混错分图告诉你为什么混。我自己的经验是厨余和其他这两类永远是重灾区因为「用过的纸巾」「沾油的塑料盒」这种边界样本人工标注都不一致。与其死磕模型不如回头把这类样本单独拎出来要么统一标注标准要么在训练时给它们更高的损失权重。做完这一步你的课程设计就不只是「调了个包」而是有数据洞察的项目答辩时老师问「你觉得哪里还能改进」你就有话说了。最后一个习惯每次改完参数把val_acc、lr、batch_size、是否解冻骨干记到一个experiments.md里。课程设计周期短很容易改着改着忘了哪组参数最好回头想复现都找不到。这个习惯我保持到现在希望帮到你。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进