简介:这份资源是一套面向高校毕设与课程作业场景的智能垃圾分类系统完整源码,适合人工智能初学者、计算机相关专业学生及需要完成类似课题的研究者。项目以图像识别为核心,涵盖数据预处理、模型构建、训练评估与可视化等环节,可作为理解CNN、PyTorch等主流技术与垃圾分类落地应用的实战范例。压缩包共52个文件,以28个Python脚本为主,覆盖数据集处理、模型定义、训练预测及可视化等模块,另含13个编译后的pyc文件与6张示意图,其余为辅的说明文档、配置及动态演示文件,整体仅1.15MB,轻量易部署。目前已有142人学习下载。通过研读源码目录与关键脚本,读者可以掌握从图片预处理到模型调优的完整流程,并借鉴其日志记录、进度可视化等工程化实现思路,为同类环保AI项目提供可复用的代码基础与排错参考。
1. 毕设里的“智能垃圾分类系统”,到底是一份什么样的工程包
如果手上这份“毕设&课程作业_用于人工智能学习所实验的智能垃圾分类系统.zip”是你接下来要交差的唯一依据,先别急着解压跑代码。人工智能方向的大作业和毕设里,垃圾分类是最常见的落地课题之一:任务边界清楚、图片数据好找、能完整覆盖数据集整理、模型训练、效果评估和界面演示这几个环节。这个zip包通常包含的是一整套可直接运行的工程——训练脚本、推理脚本、按类别组织好的图片、已经训练出的模型权重,而不是一份单纯的论文或者PPT。适合两类人:刚学完深度学习、想用真实项目验证自己能力的学生,以及需要在答辩现场完整演示“怎么训练、怎么识别、哪里会失败”的开发者。后面所有内容都围绕一套动作展开:拆包看结构,把流程跑通,再按自己的数据调参。
2. 解剖工程包:数据流、模型链路与评分规则
zip包的外壳并不重要,重要的是解压之后你看到了什么。智能垃圾分类系统的本质是一个图像分类任务,数据从文件夹进入,经过预处理、特征提取、分类层三层处理,最后输出每个类别的概率。判断一份工程包“能不能用来交差”,核心就看三件事:数据是不是完整、训练入口是不是能跑、推理入口是不是能出结果。整个第2章围绕这三件事展开。
2.1 典型工程结构:解压后对照这份目录树看
拿到zip包后先做一遍结构检查,不要直接双击跑train.py。一份标准的毕设级工程包,目录层级通常长这样:
smart_garbage/ ├── data/ │ ├── raw_images/ # 原始采集图片,未按类别分目录 │ ├── train/ # 按类别分好文件夹的训练集 │ │ ├── 可回收/ │ │ ├── 厨余/ │ │ ├── 有害/ │ │ └── 其他/ │ └── val/ # 验证集,结构与 train 一致 ├── weights/ │ └── resnet18_best.pth # 训练生成的最佳权重 ├── train.py # 训练入口 ├── infer.py # 单张图片推理入口 ├── requirements.txt └── README.md # 环境与使用说明我先解释这个目录树为什么这么设计:data目录是数据的唯一来源,train和val按类别分文件夹,这是PyTorch的ImageFolder和Keras的flow_from_directory都能直接识别的标准格式,不需要额外写读取逻辑。weights目录存放训练产物,有pth文件说明作者已经把模型训出来了,没有也没关系,后面自己训。train.py和infer.py是两个独立入口,前者负责从零把模型练出来,后者负责加载权重做单张图片推理,互不干扰。
解压后先看一个关键点:data/raw_images 里是不是平铺着大量图片,并且文件名带着类别关键字,比如“塑料瓶_001.jpg”“纸箱_002.jpg”。如果是,说明数据还没有按类别归到子目录,直接训练会报错或把所有图片当成一个类。如果data/train和data/val已经存在,说明数据整理这一步已经做完,可以跳过数据整理直接进入训练。这个判断直接决定了你后面是两小时跑通还是两天补数据,是我每次拿到工程包第一件做的事。
2.2 数据流与评分逻辑:为什么验证集准确率不是终点
系统的数据流是单向的:读图→预处理→模型推理→得到类别概率→映射成垃圾类别名称。训练脚本里的逻辑也一样,图片从train目录读入,经过随机翻转、归一化这类数据增强,进入ResNet或MobileNet主干网络,最后过一层全连接输出4到6个类别的分数。课程作业通常只看验证集准确率,但毕设答辩时老师更常问的是:“你为什么会选这个模型”“数据是怎么清洗的”“识别错了的图片是什么样”。这决定了你的评分上限不在准确率数字,而在你对整个链路有没有解释能力。
我一般会建议在理解代码跑通流程之外,额外做一份简单的错误样本分析:把预测错的图片单独复制到一个文件夹,数一数错在哪几个类别之间。比如“厨余垃圾”老是和“其他垃圾”混淆,大概率是这两类图片的拍摄背景和角度太接近,模型学到的不是垃圾本身,而是背景里的塑料袋纹路。这个问题不解决,验证集准确率再高,拿到真实场景一样翻车。
2.3 选型理由:为什么迁移学习是这个场景的默认答案
毕设垃圾分类系统的图片量级通常是几百到几千张,最少见的做法是从零训练一个深层CNN。原因一是数据太少,深层网络从零训必然过拟合——训练集准确率98%,验证集只有60%;原因二是训练时间太长,没有GPU的机器跑一个完整的ResNet要按天算。常见做法是用ImageNet上预训练好的ResNet18或MobileNetV3,冻结大部分权重,只把最后一层全连接替换成自己的类别数做微调。预训练模型已经学好了边缘、纹理、形状这些底层特征,垃圾分类需要的恰好就是这些——瓶子是光滑的圆柱形,纸箱表面有瓦楞纹,这些特征ImageNet里都有,迁移过来只需要少量数据就能拟合。
至于为什么不用更深的ResNet50或者更花哨的ViT,答案很朴素:几百张图片撑不起大模型的容量,ResNet50在三千张图片上微调,效果不一定比ResNet18好,反而训练时间和显存占用翻倍。这个选择不是技术越先进越好,而是“刚刚好能完成课程要求”最好。数据量不足的前提下,模型越简单,越不容易暴露过拟合问题。
3. 把系统跑通:从解压到一次完整训练的最小路径
现在进入动手环节。从拿到zip包到训练出模型,一共有四步:解压检查、数据整理、训练、推理验证。每一步都有一个明确产物,做完一步再看下一步,不要跳。
3.1 解压与路径检查:中文文件名和编码是第一个坑
zip包的文件名里带中文和“&”符号,在命令行里直接解压容易出乱码。Linux下解压建议用unzip加编码参数:
cd ~/workspace unzip -O gbk "毕设&课程作业_用于人工智能学习所实验的智能垃圾分类系统.zip" -d smart_garbage cd smart_garbage && find . -maxdepth 2 -type d | sort这里-O gbk的作用是告诉unzip,压缩包里的中文文件名是用GBK编码存储的,按这个编码解出来才不会乱码。macOS自带的unzip不支持-O参数,可以用Python的zipfile模块解压,文件名乱码的问题留到第3.2节一起修。Windows上用7-Zip或WinRAR右键解压基本不会有编码问题,但解压后如果看到文件名变成“锟斤拷”一类乱码,同样需要统一重命名。
解压后先确认三件事:data/train目录是否存在并包含类别子目录;requirements.txt里列了哪些依赖;模型权重文件weights目录下有没有pth文件。如果这三样都不全——比如没有data目录、只有一个train.py——说明这个包的作者只提交了代码,数据集因为太大没有放进去,你只能从第3.2节开始自己整理数据。
3.2 数据整理:把平铺图片按类别归入子目录
如果你的zip包里只有raw_images,里面全是“塑料瓶_001.jpg”“纸箱_002.jpg”这种命名,说明数据需要一个归类步骤。常见做法是按照文件名里的关键字,把图片复制到以类别命名的子目录里:
# organize_by_name.py # 把 raw_images 里 "塑料瓶_001.jpg" 形式的前缀作为类别,复制到 dataset/<类别>/ import shutil from pathlib import Path raw_dir = Path("data/raw_images") out_dir = Path("dataset") category_keywords = ["塑料袋", "塑料瓶", "玻璃瓶", "纸箱", "电池", "厨余"] for img in raw_dir.glob("*"): if img.suffix.lower() not in (".jpg", ".jpeg", ".png"): continue label = None for kw in category_keywords: if kw in img.stem: # 文件名里包含了类别关键字 label = kw break if label is None: print(f"[跳过] 无法归类: {img.name}") continue target = out_dir / label target.mkdir(parents=True, exist_ok=True) shutil.copy2(img, target / img.name) print("归类完成,各目录数量:") for d in sorted(out_dir.iterdir()): if d.is_dir(): print(f" {d.name}: {len(list(d.glob('*')))} 张")这段脚本的逻辑是从左到右遍历category_keywords列表,先在文件名里找到“塑料瓶”就把图片归到塑料瓶子目录。这里有一个细节:关键字顺序很重要,“塑料袋”和“塑料瓶”同时存在于列表里时,如果一张图叫“塑料袋包装的塑料瓶_001.jpg”,先匹配到“塑料袋”,类别就归错了。我一般会把更长的、更具体的关键字往前放,这样匹配更准确。用shutil.copy2而不是shutil.copy,是为了保留文件的元数据和拍摄时间信息,后面做数据清洗时能用上。
脚本跑完会打印每个类别的图片数量。这一步输出的数量信息很重要:如果某一个类别只有20张,另一个类别有500张,这就是一个典型的类别不平衡问题,第4.3节会专门处理。继续做训练集和验证集的划分:
# split_train_val.py # 先随机打乱,再按 8:2 切分训练集和验证集 import random import shutil from pathlib import Path src = Path("dataset") out_train = Path("data/train") out_val = Path("data/val") val_ratio = 0.2 random.seed(42) for cls_dir in sorted(src.iterdir()): if not cls_dir.is_dir(): continue imgs = list(cls_dir.glob("*")) random.shuffle(imgs) n_val = int(len(imgs) * val_ratio) for img in imgs[:n_val]: (out_val / cls_dir.name).mkdir(parents=True, exist_ok=True) shutil.copy2(img, out_val / cls_dir.name / img.name) for img in imgs[n_val:]: (out_train / cls_dir.name).mkdir(parents=True, exist_ok=True) shutil.copy2(img, out_train / cls_dir.name / img.name) print(f"{cls_dir.name}: 共{len(imgs)}张, train {len(imgs)-n_val}, val {n_val}")先把所有图片随机打乱再切分,这一步必须写在前面。如果直接按文件系统的顺序切分,同一个文件夹下连续编号的图片内容高度相似,训练集末尾和验证集开头很可能来自同一段连续拍摄,等于变相把同一批图片同时放进了训练集和验证集,验证集准确率会虚高。random.seed(42)是为了让每次切分结果一致,方便复现。
3.3 训练最小脚本:ResNet18微调,20轮内出结果
数据就绪后,需要一个能直接跑的训练脚本。我用PyTorch写一个ResNet18迁移学习的完整最小版本,这是毕设垃圾分类最常见、也最不容易出问题的组合:
# train.py # 用 ImageNet 预训练权重初始化模型,替换最后一层做微调 import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, models, transforms device = "cuda" if torch.cuda.is_available() else "cpu" num_classes = 4 batch_size = 32 epochs = 20 learning_rate = 1e-3 train_tf = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p=0.5), 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("data/train", train_tf) val_ds = datasets.ImageFolder("data/val", val_tf) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2) val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=2) model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) model.fc = nn.Linear(model.fc.in_features, num_classes) model = model.to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5) best_acc = 0.0 for epoch in range(epochs): model.train() total_loss, correct, total = 0.0, 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() total_loss += loss.item() * images.size(0) correct += (outputs.argmax(1) == labels).sum().item() total += images.size(0) scheduler.step() model.eval() val_correct, val_total = 0, 0 with torch.no_grad(): for images, labels in val_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) val_correct += (outputs.argmax(1) == labels).sum().item() val_total += images.size(0) train_acc = correct / total val_acc = val_correct / val_total print(f"Epoch {epoch+1:02d} | loss {total_loss/total:.4f} | train_acc {train_acc:.4f} | val_acc {val_acc:.4f}") if val_acc > best_acc: best_acc = val_acc torch.save(model.state_dict(), "weights/resnet18_best.pth") print(f" 保存最佳模型,val_acc = {val_acc:.4f}") print("训练完成,最佳验证准确率: {:.4f}".format(best_acc))先说模型初始化这行。新版本torchvision里,models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)是官方推荐的写法,如果你在别的教程里看到models.resnet18(pretrained=True),那是旧写法,新版已经会触发警告甚至直接报错,这事在避坑章节还会重点讲。替换最后一层全连接,是因为ImageNet预训练模型输出的1000类结果对我们没用,改成4类之后只训练这一层和前面层的微调参数。
参数设计上,learning_rate=1e-3是微调场景的通用起点,Adam优化器不需要额外调momentum。StepLR每5轮把学习率减半,作用是让模型在训练后期用小步长收敛得更稳。train和val用不同的预处理流程:训练集加了RandomHorizontalFlip做数据增强,验证集不做任何随机变换,只把图片缩放和归一化,这样才能公平评估。保存模型的判断标准是验证集准确率而不是训练集准确率,每次都存最高的那一次。
3.4 单张图片推理:没有CUDA也要能演示
训练完成后,还需要一个独立的推理脚本,这是答辩时演示用的核心。它要能在没有GPU的机器上运行,并且对单个图片文件直接给出预测结果:
# infer.py # 加载训练好的权重,对单张图片输出类别与置信度 import sys import torch from PIL import Image from torchvision import models, transforms device = "cuda" if torch.cuda.is_available() else "cpu" num_classes = 4 label_names = ["可回收", "厨余", "有害", "其他"] # 顺序和训练时 ImageFolder 的目录一致 model = models.resnet18(weights=None) model.fc = torch.nn.Linear(model.fc.in_features, num_classes) model.load_state_dict(torch.load("weights/resnet18_best.pth", map_location=device)) model.to(device).eval() infer_tf = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) def predict(img_path: str) -> tuple[str, float]: img = Image.open(img_path).convert("RGB") tensor = infer_tf(img).unsqueeze(0).to(device) with torch.no_grad(): probs = torch.softmax(model(tensor), dim=1)[0] idx = int(probs.argmax()) return label_names[idx], probs[idx].item() if __name__ == "__main__": label, confidence = predict(sys.argv[1]) print(f"预测类别: {label},置信度: {confidence:.2f}")推理脚本最容易忽略的一行是Image.open(img_path).convert("RGB")。如果传入的图片是RGBA格式的PNG或者灰度图,convert成RGB能保证后续处理的张量形状始终是3通道,否则会在这里报通道数不匹配的错。load_state_dict时加了map_location=device,这样在CPU机器上加载GPU训练的权重也不会报错。推理时用torch.no_grad()包裹,不保留计算图,显存占用和内存占用都会大幅下降。
运行方式很简单:python infer.py 某张图片的路径.jpeg。输出的置信度是Softmax之后的最大概率值,和训练时的类别索引一一对应。这里要特别说明label_names的顺序必须和训练时ImageFolder自动生成的类别顺序一致,ImageFolder的类别排序默认按照文件夹名称的字母序,如果训练脚本里的顺序是“可回收、厨余、有害、其他”,推理脚本里的label_names也要严格按这个顺序写,否则预测标签会张冠李戴。
4. 调参指南:学习率、batch_size与模型选型的边界
训练脚本能跑通只是第一步,怎么让模型在几百张图片上达到答辩能看的准确率,才是花时间的地方。这一章不聊玄学,只给一套有依据的参数起点和调整逻辑。
4.1 一套能用的起点参数
我直接给一个经过多次验证的起点参数表,适合大部分毕设级别的垃圾分类数据集:
| 参数 | 推荐起点 | 说明 |
|---|---|---|
| learning_rate | 1e-3 | 微调预训练模型的通用起点,loss震荡时降到1e-4 |
| batch_size | 32 | 显存不够降到16或8,同时学习率可以同步降低 |
| epochs | 15~20 | 几百张图片20轮内基本收敛,再多轮次收益很小 |
| 输入尺寸 | 224x224 | ResNet/MobileNet的默认输入,改成512需要重新设计网络 |
| optimizer | Adam | 默认参数即可,简单稳定,不需要额外调momentum |
| 学习率衰减 | StepLR,step_size=5,gamma=0.5 | 每5轮减半,让后期收敛更慢更稳 |
这里最容易被忽视的是batch_size和学习率的联动关系。batch_size翻倍时,梯度估计更准,可以适当调大学习率;batch_size减半时,学习率也应该跟着减半。很多翻车场景是显存不够强行把batch_size从32降到8,学习率仍然是1e-3,结果训练时loss一路震荡不下降。这不是模型问题,是超参数没有联动调整。
epochs的设置要结合图像数量看。500张训练图,ResNet18微调,10到15轮之内验证准确率就会趋于平稳。如果20轮后准确率还在缓慢上升,可以再加10轮;如果20轮时训练准确率接近100%而验证准确率停滞不动,那是过拟合信号,加轮次没有用,要回头查数据有没有问题或者加数据增强。
4.2 ResNet18、MobileNetV3、EfficientNet怎么选
毕设演示的设备和训练设备往往不是同一台,很多人是在有GPU的实验室训练,到答辩现场用一台没有独立显卡的笔记本做演示。模型选型时不能只盯着训练精度,还要看推理速度是否扛得住现场演示。
| 模型 | 参数量 | CPU单张推理耗时 | 适用场景 |
|---|---|---|---|
| ResNet18 | 约11M | 200~400ms | 默认首选,训练快、资料多、最稳 |
| MobileNetV3-Small | 约2.5M | 80~150ms | 演示机没有GPU,或需要实时预览 |
| EfficientNet-B0 | 约5.3M | 300~500ms | 数据量少但追求精度上限时 |
我的建议是训练阶段直接用ResNet18,不需要纠结。如果答辩现场用CPU机器演示,再把训练好的权重迁移到MobileNetV3上做一次微调,或者干脆一开始就训练MobileNetV3。MobileNetV3的准确率在几百张垃圾分类图上比ResNet18低一到两个点,但CPU推理速度快一倍,现场体验好很多。
EfficientNet在这个任务里有点尴尬。它理论上的精度上限更高,但在几百张图片的小数据集上优势不明显,训练时间反而更长,调试成本也更高。课程作业和毕设的评分标准里几乎不会因为你用了EfficientNet而多给分,却很可能因为演示卡顿扣分。模型选型的边界在这里:够用就好,把省下来的时间花在数据清洗和错误样本分析上,得分更高。
4.3 类别不平衡:准确率虚高的真相
几百张数据里最常见的坑是类别分布极不均匀。比如“可回收”500张,“有害”却只有30张。如果直接训,模型会靠猜“可回收”拿到95%左右的验证准确率——因为“有害”样本太少,模型基本不会预测这个类别,但准确率数字看起来还很漂亮。答辩时老师随便拿一节废电池测,模型就会露馅。
有两个常用解决方案,第一个是加权采样,让少数类样本在每一个batch里出现的概率更高:
# weight_sampler.py # 给样本数少的类别分配更高的采样权重 from pathlib import Path from torch.utils.data import WeightedRandomSampler class_names = ["可回收", "厨余", "有害", "其他"] cat_dir = Path("data/train") # 统计每个类别的样本数,样本越少权重越高 class_counts = [len(list((cat_dir / name).glob("*"))) for name in class_names] weights = [1.0 / cnt for cnt in class_counts] sample_weights = [] for _, label in train_ds: # train_ds 是第3.3节里创建的 ImageFolder sample_weights.append(weights[label]) sampler = WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), replacement=True) train_loader = DataLoader(train_ds, batch_size=32, sampler=sampler)这段代码与第3.3节训练脚本的差别只在DataLoader这一步。把原来的shuffle=True去掉,改用sampler,每个batch里少数类出现的次数就会明显增多。WeightedRandomSampler的核心参数是sample_weights,它的长度必须和数据集样本总数一致。replacement=True表示每次采样可以重复抽取同一个样本,这在小数据集上是必须的,否则少数类会被很快抽完。
第二个方案是数据增强,对少数类做更多的随机变换模拟出更多样貌:把RandomHorizontalFlip换成RandomRotation、ColorJitter的组合。但数据增强只能治标,样本数量差异超过20倍时,最可靠的方案还是扩充数据。找数据集的时候注意,公开的垃圾分类数据集里“有害垃圾”这个类别通常就是最少的,因为收集成本高。所以拿到一个新工程包,第一件事统计各目录图片数量,这个习惯能提前暴露后面80%的准确率问题。
5. 智能垃圾分类系统必踩的5个坑:现象、原因与修复办法
这一章直接给踩坑记录,每一条都是“现象→原因→解决”,照着排查就行。这里没有玄学,只有经验。
5.1 坑1:torchvision的pretrained参数写法翻车
现象:训练脚本一启动就报TypeError: resnet18() got an unexpected keyword argument 'pretrained',或者只弹出FutureWarning但不报错,准确率却始终很低。
原因:torchvision 0.13版本之后,官方把models.resnet18(pretrained=True)标记为弃用,推荐用weights参数。不同版本之间pretrained参数的表现不一致——有的直接报错,有的还能跑但打印警告,有的静默失败导致模型根本没有加载预训练权重,从随机初始化开始训练,几百张图片根本训不动。
解决:把初始化那行改成models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)。如果代码里复用了很多网上找的旧版教程,全局搜pretrained=True全部替换掉。这修正不了模型精度低的问题,因为前提是权重根本没对上。
5.2 坑2:中文文件名在Windows下乱码
现象:zip包在Linux下解压正常,拷到Windows解压后train目录下的“可回收”变成“鍙?敹鍥炴敹”这一串乱码,ImageFolder读不到任何类别目录,训练直接空转。
原因:zip包内中文文件名用了GBK编码,而Windows解压工具的默认编码在不同系统区域设置下不一致,导致文件名解码错乱。
解决:如果尚未解压,Linux下用unzip -O gbk,macOS下用Python的zipfile模块解压后用4.1节的关键字归类脚本重新整理数据目录。如果已经解压坏了,最省事的办法是在Windows里用7-Zip重新解压一次,并把系统区域设置里“Beta版使用Unicode UTF-8提供全球语言支持”打开,然后重新运行一次数据整理脚本。这个坑最隐蔽的是脚本能运行但数据集目录全是乱码,训练出的模型类别名也都是乱码,答辩时根本没法解释。
5.3 坑3:验证集划分时数据泄漏,准确率失真
现象:训练集准确率80%,验证集准确率却高达97%,两者差距明显不合理。进一步查看验证集样本,发现很多图片和训练集里某几张几乎一样。
原因:数据划分前没有做随机打乱。图片是按采集时间顺序编号的,连续编号的图片内容几乎相同——同一堆垃圾从不同角度连拍了几十张,或者同一段视频每隔10帧截一次图。按文件系统顺序直接切8:2,训练集的末尾和验证集的开头来自同一批连续拍摄,等于用见过的图去考模型。
解决:划分之前先random.shuffle,这一点我在3.2节的代码里已经写进去了。如果数据是视频抽帧来的,更稳妥的做法是按视频分组划分:同一个视频产生的帧全部进训练集或全部进验证集,尽量避免跨集合泄漏。鉴定有没有数据泄漏最快的方法,是从验证集里随机抽一批图,肉眼和训练集对比,发现成片重复就直接重做划分。
5.4 坑4:演示现场显存崩溃或推理太慢
现象:训练和推理都在自己的电脑上没问题,换到答辩现场的低配笔记本上,单张图片推理等了5秒才出结果,现场气氛尴尬。
原因:现场机器可能只有CPU没有独立显卡,而模型是ResNet18,CPU推理一张224x224的图片本来就要200到400毫秒。如果infer脚本里忘记加torch.no_grad(),或者推理时把整个验证集都加载了一遍,时间会翻倍。另外,如果模型用了ResNet50这类更大网络,CPU推理延迟会直接冲到1秒以上。
解决:推理脚本固定用no_grad()包裹,确认没有梯度计算开销。演示机没有GPU时,把模型换成MobileNetV3重新做一次微调,CPU推理延迟能降到100毫秒左右。如果不想重新训练,另一个方案是把模型导出成ONNX再用ONNX Runtime推理,这部分在第6.2节展开。总之,演示机器的算力边界要提前确认,不要到现场才测。
5.5 坑5:输入图片通道不对导致推理报错
现象:训练一切正常,推理时部分图片报错“Expected 3 input channels, got 4”或者“Expected 3 input channels, got 1”。
原因:手机或相机拍出的PNG图片可能是RGBA四通道,一些扫描件或老照片直接是灰度图单通道,ResNet的卷积层要求输入必须固定为3通道RGB。训练脚本里数据已经统一处理过,推理脚本里如果直接Image.open不做转换,就踩这个坑。
解决:推理脚本里统一加一句.convert("RGB"),这也是3.4节代码里特意写上的原因。另外,训练脚本的预处理里也可以加一个转换,保证不管谁来训练,数据输入始终是统一格式。灰度图转RGB不会增加有效信息,但能保证流程不断。
6. 验收与进阶:用“没见过的图片”证明系统真的能用
6.1 用混淆矩阵和新照片做交叉验证
训练完成之后,先别急着写答辩PPT。我建议做两件事,第一件是输出混淆矩阵和分类报告,把所有类别的精度、召回率摊开看:
# evaluate.py # 在验证集上输出分类报告与混淆矩阵 from sklearn.metrics import classification_report, confusion_matrix import numpy as np y_true, y_pred = [], [] model.eval() with torch.no_grad(): for images, labels in val_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) y_true.extend(labels.cpu().numpy()) y_pred.extend(outputs.argmax(1).cpu().numpy()) print(classification_report(y_true, y_pred, target_names=label_names)) print(confusion_matrix(y_true, y_pred))分类报告里的precision和recall比总体准确率更有参考价值。如果“有害”类别的recall是0.2,说明100张有害垃圾图片只认出了20张,这个类别必须处理。第二件事是手机拍几张真实生活垃圾图片,新拍的、完全没参加过训练的,放进一个独立文件夹跑推理。这个验证方式比测试集准确率更有说服力——答辩时老师当场掏出手机拍一张废纸箱,你能当场给出正确结果,这个印象分比任何指标都重要。
6.2 导出ONNX:把系统从PyTorch挪到CPU推理
如果答辩机没有PyTorch环境,或者不想现场装一堆依赖,可以把模型导出成ONNX格式,用ONNX Runtime做CPU推理:
# export_onnx.py # 把训练好的 PyTorch 模型导出为 ONNX 格式,便于跨平台部署 dummy = torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy, "garbage_cls.onnx", input_names=["image"], output_names=["logits"], dynamic_axes={"image": {0: "batch"}}, opset_version=12, do_constant_folding=True, )导出后可以直接用onnxruntime在CPU上做推理,不需要PyTorch,不需要显卡,依赖只剩onnxruntime和numpy。dynamic_axes参数把batch维声明为动态,这样导出后的模型既能推理单张图片,也能一次推理多张,灵活性更好。opset_version=12是兼容性很好的选择,新版旧版的onnxruntime都能加载。这一步做完,系统就从“一个PyTorch训练实验”变成了“一个能部署的模型服务”,这在毕设评分里是一个明确的加分项。
进阶方向再提醒一句:如果课程要求做界面,优先考虑用Flask起一个本地网页而不是Tkinter,网页演示不需要对方装Python环境,视觉效果也好得多;但无论界面怎么做,后端必须统一走infer.py里这个predict函数,端口怎么开、界面怎么摆都不影响模型本身的逻辑。
我做这个项目时养成的习惯是:每次训练完都会手动保留一份“错误样本文件夹”,把验证集里预测错的图片全部复制进去,截图留在答辩PPT里。这样做不是为了证明模型完美,而是为了说明你知道模型在哪里会犯错、为什么犯错、下一步怎么改。老师问“你这个系统还有什么不足”时,能针对具体图片说出改进方向的学生,比只背指标的学生拿分高得多。这套流程从数据整理到模型导出,走两遍之后,你会发现后面再做任何图像分类项目,都能直接复用同一条链路。希望帮到你。
本文还有配套的精品资源,点击获取