ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

用Python深度学习实现垃圾分类图像分类:从迁移学习到部署

用Python深度学习实现垃圾分类图像分类:从迁移学习到部署 简介这是一套基于深度学习的垃圾分类小工程压缩包面向具备一定Python基础、希望快速上手图像分类实战的开发者与学生。项目以PyTorch和ResNet50为核心包含数据索引、验证集划分、数据生成脚本、模型定义、训练脚本、演示脚本及理论笔记围绕垃圾分类场景串联起从数据准备、模型训练、评估到推理演示的完整流程。压缩包共14个文件以Python脚本为主体配合中英文Markdown说明、Word文档中间理论、训练信息、JSON配置与示例图片其中resnet50.py定义网络结构trainNet.py执行训练Demo.py用于推理演示genit.py负责数据生成结构清晰且易于二次开发。整体压缩包仅4.22MB轻量便捷适合本地快速解压查看。目前已有1210人学习使用可作为课程设计、毕业设计或AI入门小项目的参考实现帮助读者减少环境搭建和调参试错成本快速形成可运行的垃圾分类Demo。1. 为什么垃圾分类小工程值得用 Python 深度学习跑一遍垃圾分类这几年被写进简历的次数不少但真正能跑的 Python 深度学习小工程其实不多。多数试水的项目卡在一个环节公开数据集太干净自己拍的照片又太少训练出来的模型一到真实场景就失灵。反过来如果愿意把垃圾分类当成一个完整的小工程去做从数据整理、模型训练到单图预测周期可以压缩到一周以内成本低到普通电脑就能承受。“小工程”三个字在这里的含义是先把最小可用闭环做出来而不是一上来就堆目标检测之类的重方案。做这件事的人通常是两类。一是刚学完 PyTorch 基础、想用深度学习项目练手的学生想通过垃圾分类把卷积网络、迁移学习、数据增强这些概念串起来二是公司里需要快速验证垃圾分类 AI 方案能不能行的工程师更关心用最少的标注量拿到一个说得过去的准确率。他们的共同点是不想被算法复杂度拖住想先把准确率跑到一个能用的位置。做到什么程度算完我的标准是准备好分类数据集用 ResNet 或 MobileNet 完成训练模型对一张没见过的垃圾照片能输出桶类别再封装成命令行预测脚本。下面这条链路就按这个顺序展开每个环节给参数也给坑位。2. 为什么垃圾分类这类小工程用图像分类而不是目标检测2.1 业务边界决定算法选型先确认你要回答的问题垃圾分类这个需求听起来是一个识别问题但识别有两个层次。图像分类网络回答的是“这张图属于哪一类”目标检测网络回答的是“图里有哪些物体、分别在哪、是什么”。放在垃圾桶前这个场景里用户举起手机对着一个易拉罐拍一张照系统告诉他该扔进“可回收”还是“其他”。图片主体通常只有一个位置大概率在画面中央所以你真正需要的只是“类别”这个答案。目标检测在这个流程里多干的画框活对最终决策没有帮助反而要标注边界框、要处理输出分支、要调 NMS 参数成本直线上升。有一类场景确实要转目标检测比如垃圾房顶上的监控摄像头要同时识别画面里多个垃圾袋还要判断哪个桶满了。这种多目标、大画面场景才需要检测甚至跟踪。小工程最重要的一件事是把问题边界画窄。我见过好几个朋友第一版就上 YOLO花两周标注完框最后效果还不如一个轻量分类网络。判断标准很简单输入是“主体明确的单物体图片”还是“开放场景的画面”。前者选分类后者再考虑检测。这里还牵扯到模型规模。分类网络可以直接吃 224×224 的输入预训练权重满大街都是一张消费级显卡能训完推理时甚至不用 GPU检测网络最小也得跑 640×640 的输入还要做后处理。小工程的资源预算通常很紧能在分类上解决的事不要往检测上靠。这句话听起来像废话但真到了自己动手的时候很容易被“检测看起来很酷”带跑。2.2 主干网络选型ResNet18、ResNet50 还是 MobileNetV3选主干之前先确定训练策略。如果走迁移学习主干选择影响的是“提取底层特征的能力”和“训练成本”。我整理过一张对照表数值不是绝对基准而是我在不同机器上跑出来的体感范围主干参数量224×224 单张 CPU 推理8G 显存训练体验适合场景ResNet18约 11.7M百毫秒级还行很轻松batch 64 也能跑小工程首选ResNet50约 25.6M明显慢一截只能小 batch显存紧张数据量大、追求极限精度MobileNetV3-Small约 2.5M几十毫秒级最快很轻松要部署到手机或树莓派我一般从 ResNet18 开始。理由不是它精度最高而是它的调试体感最平滑学习率稍微给大一点不会立刻发散训练曲线比较典型网上同架构的踩坑讨论最多。ResNet50 在小数据上带来的精度提升通常只有一到两个点但训练时间和显存压力实打实。MobileNetV3 的优势要等后续做剪枝、量化才能放大如果只是课程设计或者方案验证先跑通流程更重要。选型还有一个顺手的小技巧无论选哪个主干直接用 torchvision 里带预训练权重的版本。自己从头训练的 ResNet18 在几千张垃圾图片上收敛很慢精度也难看用 ImageNet 预训练权重可以少走很多弯路。迁移学习的起点和边界就是下一节要说的内容。2.3 迁移学习不是“拿权重就完事”而是“先冻结再微调”垃圾图片和 ImageNet 的图片差异其实不小。ImageNet 里是狗、猫、汽车、自然场景垃圾数据集里多是形状不规则的瓶瓶罐罐、皱巴巴的纸、沾了污渍的塑料。但底层卷积核提取的边缘、纹理、颜色渐变这些通用特征在两个域之间是能迁移的所以迁移学习一定比从零训练快得多。这也是小工程能把训练周期压到一周内的根本原因。我的默认做法是加载在 ImageNet 上训练好的 ResNet18替换最后的全连接层。原来的 fc 层输出 1000 类现在换成 6 类。第一次跑的时候把整个网络除 fc 以外的参数全部冻结只用小学习率训练新的分类头3 到 5 个 epoch让新层先适应数据分布。然后解冻最后一个残差块也就是 layer4用更小的学习率做整体微调。两步顺序别反否则新初始化的层会输出很大的梯度把预训练权重冲乱loss 很容易在训练初期就飙升。import torch.nn as nn from torchvision import models # 新版 torchvision 用 weights 参数加载预训练权重 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) in_features model.fc.in_features model.fc nn.Linear(in_features, 6) # 换成垃圾类别数 # 先冻结除分类头外的所有参数 for name, param in model.named_parameters(): if name.startswith(fc): param.requires_grad True # 分类头不冻结 else: param.requires_grad False这段代码的关键是 requires_grad 开关。第一阶段优化器只接收 fc 层的参数其他层的梯度不计算节省显存也避免扰动。到了微调阶段再把 layer4 的参数打开设置成不同学习率。这里的“先冻结再微调”顺序直接影响 loss 曲线是平滑还是来回蹦跳。第一阶段的全连接层不提取特征只负责把特征映射到 6 个类别上第二阶段解冻 layer4 是为了让高层语义特征更贴近垃圾图片的分布但学习率必须降下来不然预训练权重会被新数据改得面目全非。3. 数据整理与预处理把垃圾图片文件夹变成可训练的 Dataset3.1 先定目录结构别在数据形态上翻车垃圾分类小工程最常见的数据集形态是 train、val、test 三个大目录每个类别各占一个子目录。这种结构是 torchvision.datasets.ImageFolder 的标准输入几乎不用写自定义读取代码。公开数据集方面英文资料里常看到 TrashNet六个类是 cardboard、glass、metal、paper、plastic、trash国内平台则常见四十多类的中文垃圾分类数据集把厨余、可回收、有害、其他拆得很细。第一次动手不要贪多先选 6 到 8 个类把链路跑通后面再加。类别数变多时模型最后的 fc 层维度要跟着变训练成本也会上升。在拿到数据集之后我习惯先写一小段统计脚本看看每个类到底有多少张图。这一步很多人会跳过但其实价值很高它能把“类别不均衡”这个雷提前排掉。garbage_6/ ├── train/ │ ├── cardboard/ │ ├── glass/ │ ├── metal/ │ ├── paper/ │ ├── plastic/ │ └── trash/ ├── val/ └── test/import os from collections import Counter root garbage_6 for split in [train, val, test]: split_dir os.path.join(root, split) class_names [d for d in os.listdir(split_dir) if os.path.isdir(os.path.join(split_dir, d))] counter Counter() for c in class_names: c_dir os.path.join(split_dir, c) counter[c] len(os.listdir(c_dir)) print(split, dict(counter))这段代码统计每个子文件夹的图片数量主要用来发现两类问题一是某个类别图片数量和别的类差了一个数量级二是 val、test 的分布和 train 不一致。小工程因为数据量小很容易出现 train 里某个类 800 张、val 里只有 10 张的情况最后算出来的指标就不可信。这里还有一个很少被提到的细节ImageFolder 对类别顺序依赖文件系统的排序所以目录名如果用中文后续画混淆矩阵、打印日志时容易出现编码问题。我习惯把目录名转成英文别名或拼音比如 glass、metal、plastic而不是“玻璃”“金属”“塑料”。这个习惯帮我少处理了很多莫名其妙的 Windows 编码报错。3.2 ImageFolder 好用但 DataLoader 的参数别照抄用 ImageFolder 构建数据集的时候要先把训练集和验证集的 transform 区分开。训练集做随机增强验证集只做缩放和标准化不引入随机性否则评估指标会上下抖动看不出真实水平。from torch.utils.data import DataLoader from torchvision import datasets, transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees10), transforms.ColorJitter(brightness0.2, contrast0.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]), ]) train_dataset datasets.ImageFolder(rootgarbage_6/train, transformtrain_transform) val_dataset datasets.ImageFolder(rootgarbage_6/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers2, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse)这个代码块里的参数都是有讲究的。Resize((224, 224)) 是分类网络的常驻输入尺寸我用 224 而不是 192 或 320是为了匹配 ImageNet 预训练权重训练时的分辨率习惯。RandomHorizontalFlip 对垃圾照片是安全的因为瓶子横着竖着都不改变类别RandomRotation 我只敢给到 10 度再大就会把易拉罐和塑料瓶的长宽比破坏掉。ColorJitter 的 brightness 和 contrast 都是 0.2用来模拟不同灯光下的拍摄效果。DataLoader 里值得单独说的是 num_workers 和 pin_memory。num_workers 在 Linux 上可以给 4 或 8在 Windows 上经常变成灾难这一点在第 5 章会单独讲。pin_memoryTrue 在 GPU 训练时能减少数据从 CPU 拷贝到显存的时间但如果你用的是 CPU 训练这个参数不生效可以去掉。增强还有一个误区有人喜欢把 RandomResizedCrop 也加进去用随机裁剪模拟不同视角。这对狗、猫这类“中心物体”合理但对垃圾图片来说太危险一个随机区域很可能把瓶子从中间截断网络只看到半截类别信息就没了。我一般不加加了反而会拉低准确率。3.3 类别不均衡先加权再谈模型能力垃圾分类数据集的另一个常客是类别不均衡。比如“其他垃圾”可能有两千张“有害垃圾”只有几十张因为有害垃圾本来就少。这种情况下模型会偷懒只要把所有图片都判成“其他”整体准确率也能轻松上 70%但对你没有任何意义。常见有两种解法。第一种在数据采样层面做用 WeightedRandomSampler 让少数类被采到的概率更大本质上是过采样第二种在损失函数层面做把 CrossEntropyLoss 的 weight 参数设成各类样本数的倒数让少数类的错分贡献更大。小工程我更推荐第二种改动最小一个参数就能做到。import torch.nn as nn import numpy as np class_counts np.array([1200, 80, 600, 900, 1100, 150]) # 每类样本数 weights 1.0 / class_counts weights weights / weights.sum() criterion nn.CrossEntropyLoss(weighttorch.from_numpy(weights).float())这里 weight 数组的顺序必须和 ImageFolder 的类别顺序一一对应。ImageFolder 默认按目录名字母顺序排列所以如果你的类别列表是 [cardboard, glass, metal, paper, plastic, trash]那 class_counts 也必须是同样的顺序。顺序一旦错位模型会学出一套“错位逻辑”训练集 loss 降得下去但验证结果乱成一锅粥。4. 训练实验的两次跑通先训分类头再微调主干4.1 完整的两段式最小训练脚本到了训练这一步最忌讳的是把网上大模型的训练代码原样搬过来。垃圾分类单分类任务模型小、数据少脚本没必要写成一两百行。我一般把训练脚本分成两个阶段第一段冻结主干只训 fc第二段解冻 layer4 做微调。下面这个脚本去掉了早停、日志等装饰保留核心逻辑方便照着改成自己的版本。import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, models, transforms device torch.device(cuda if torch.cuda.is_available() else cpu) num_classes 6 freeze_epochs 5 finetune_epochs 15 train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers2) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse) model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, num_classes) model model.to(device) # 第一阶段只优化分类头 for p in model.parameters(): p.requires_grad False for p in model.fc.parameters(): p.requires_grad True optimizer torch.optim.Adam(model.fc.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() model.train() for epoch in range(freeze_epochs): for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() loss criterion(model(images), labels) loss.backward() optimizer.step() print(ffreeze epoch {epoch1}, loss {loss.item():.4f}) # 第二阶段解冻 layer4用更小学习率微调 for p in model.layer4.parameters(): p.requires_grad True optimizer torch.optim.SGD([ {params: model.fc.parameters(), lr: 1e-3}, {params: model.layer4.parameters(), lr: 1e-4}, ], momentum0.9) model.train() for epoch in range(finetune_epochs): for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() loss criterion(model(images), labels) loss.backward() optimizer.step() print(ffinetune epoch {epoch1}, loss {loss.item():.4f})这个脚本能跑通的核心是“两段式”。第一阶段用 Adam 专门训新加的分类头因为新的 fc 层权重是随机初始化的梯度幅度比较大需要单独喂一个小学习率如果把整个网络都解开预训练权重会被新层的梯度扰动前几个 epoch 的 loss 曲线就会像心电图的颤振。第二阶段换成 SGD 加动量是因为微调阶段需要更稳的权重更新方向Adam 的自适应学习率在这种情况下反而容易在小数据集上过拟合。第二阶段的参数列表也用得很刻意fc 层的学习率 1e-3layer4 是 1e-4两者差了一个数量级。原因是越靠近输入的层预训练权重越通用越不需要改动越靠近输出的层越需要迎合当前数据集。layer4 和 fc 之间这种层级差异就体现在学习率的差异上。4.2 关键参数怎么定学习率、batch size、epoch 与优化器参数这个东西在小工程里最容易被当成玄学。实际上它是有边界可循的。下面是这套六分类任务我常用的初始区间数据量不大时基本不用改太多参数建议值说明第一阶段优化器Adam对随机初始化的 fc 层收敛快第一阶段学习率1e-3再大容易把 fc 梯度带崩第二阶段优化器SGD(momentum0.9)微调阶段更稳泛化更好第二阶段 fc 学习率1e-3分类头已经训练过的可以保持第二阶段 layer4 学习率1e-4解冻层改动不动降低扰动batch_size328G 显存能放下4G 显存可以减到 16freeze_epochs3~5太久没意义fc 层很容易过拟合finetune_epochs10~20边训练边看验证 loss设早停更好关于 batch size有一个经常被忽略的影响batch 越大梯度方向越平滑但显存压力越大batch 太小每个 step 的梯度噪声大训练曲线会抖。我一般先在 8G 显存上试 batch 32如果显存还有余量就加到 64但小数据集没必要硬撑32 在收敛速度和稳定性上通常就够了。训练时还有一个小技巧设置随机种子。PyTorch 的 DataLoader 在多进程场景下如果不固定 seed两次训练得到的验证结果会有几个点的波动。我会在脚本开头加上 random.seed(42)、torch.manual_seed(42)并且在多卡或分布式场景里再补一句 torch.cuda.manual_seed_all(42)。这个小习惯成本极低但对复现结果帮助很大。4.3 验证用混淆矩阵代替“看准确率”训练完之后很多人只看整体准确率就收工。实际上垃圾分类这种多分类任务整体准确率会掩盖很多问题。比如有害垃圾只占 5% 的样本模型把所有有害垃圾都错判成了其他垃圾整体准确率照样有 95%。所以我的习惯是每次训练完都生成混淆矩阵按类别看召回率。import numpy as np from sklearn.metrics import confusion_matrix def evaluate(model, loader): model.eval() all_preds, all_labels [], [] with torch.no_grad(): for images, labels in loader: images, labels images.to(device), labels.to(device) preds model(images).argmax(dim1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) return cm, all_labels, all_preds cm, _, _ evaluate(model, val_loader) print(cm) # 按类别打印召回率 for i, name in enumerate(class_names): recall cm[i, i] / cm[i].sum() if cm[i].sum() 0 else 0 print(f{name}: {recall:.2%})这段代码输出的是一个 6×6 的矩阵对角线上的数字越大越好。如果某一行非对角线有大值比如 glass 那一行的数字跑到 plastic 那一列就说明这两类在特征空间里挨得很近模型分不清。看到这种结果下一步不是盲目加数据而是要去看样本本身是不是训练集里的玻璃瓶和塑料瓶本来就长得像或者拍摄背景一致导致网络根本没学到区分性特征。5. 垃圾分类训练里最常遇到的5个坑现象、原因、解法这一章的素材来自我实际跑垃圾分类时踩过的坑也在学员群里看到不少人反复掉进去。每个坑按“现象、原因、解决”拆开写目的很简单就是让你在遇到类似症状时能快速定位。5.1 训练集准确率很高真实拍摄照片一测就翻车现象在划分好的验证集上准确率到了 93%拿手机对着桌面上的瓶子拍一张预测结果却是错的甚至置信度还挺高。原因公开数据集的图片大多拍摄条件单一背景干净、光照均匀、物体完整居中真实世界里的垃圾照片是厨房灯光、随手拍、遮挡、倾斜、反光分布差异很大。这就是数据分布偏移模型在“干净域”上学到的特征到了“脏域”就失效了。解决最有效的手段是把真实照片拉进验证流程。我建议在项目里建一个 real_test 文件夹里面放 20 到 30 张手机拍摄的垃圾照片每次训练完跑一遍用这个结果判断能不能交付而不是只盯着 val 准确率。如果真实照片准确率太低就挑一部分放回训练集再配合第 3 章的增强参数慢慢拉齐分布。这个 real_test 文件夹我到现在还在用它就是模型的“后悔药”。5.2 玻璃和塑料这两类总在互相误判现象混淆矩阵里glass 那一行非对角线的值大部分落在 plastic 列反过来也一样。单独看这两个类的召回率可能都在 70% 上下。原因透明玻璃瓶和透明塑料瓶在白色背景、统一灯光的拍摄条件下外观实在太接近了。分类网络主要靠纹理、颜色、形状做判断而这两类在这三个维度上重叠度都很高。解决回到业务上看这两个类是不是真的必须分开。如果你的目标是“可回收”和“其他”那玻璃和塑料完全可以合成一个“可回收”类模型压力立刻减半。如果业务上必须分开那就需要在样本层面补信息拍摄时带上瓶盖、标签、包装等细节或者用更高分辨率的输入让网络看到瓶口、底部的差别。还有一个偏门思路是换成更细的模型输入让网络对高光区域更敏感但这已经超出小工程的范围了。5.3 loss 卡在 1.79 附近不动准确率像随机猜现象训练刚开始 loss 从 1.79 附近起步十几个 epoch 下来纹丝不动准确率徘徊在 16% 左右。原因六个类别的随机猜测 loss 约等于 ln(6)≈1.79模型完全没在学习。最常见的原因是 Normalize 写错或者漏填比如把 mean 全设成 0、std 全设成 1输入分布和预训练权重期待的对不上其次是标签顺序错位ImageFolder 的类别顺序和你打印出来的类别列表不一致。解决第一时间做两件事。一是把数据增强里的 Normalize 参数改回 [0.485, 0.456, 0.406] 和 [0.229, 0.224, 0.225]别自己发明新数字二是打印一个 batch 的标签对照图片内容确认类别映射。这两步都检查完再试着把学习率从 1e-3 提到 3e-3看 loss 有没有动静。还没动静就检查 DataLoader 的 shuffle 是否开启顺序喂数据会让模型学到记忆而不是规则。5.4 Windows 下 DataLoader 频繁卡死或报 BrokenPipeError现象在 Windows 上跑训练脚本num_workers2每跑几个 epoch 就卡住或者直接报 BrokenPipeError重启内核也没用。原因Windows 下 PyTorch 的 DataLoader 多进程使用了和 Linux 不同的进程创建方式加上如果代码跑在 Jupyter Notebook 里进程管理更容易出问题。这是 Windows 上深度学习环境里出了名难缠的坑和你 Python 装的是 3.8 还是 3.11 关系不大。解决最快的临时方案是把 num_workers 改成 0让数据加载在主进程里同步完成训练速度会慢一点但稳定性立竿见影。如果要在 Windows 上保留多进程就把训练代码整体放进if __name__ __main__:下面避免在 import 时反复触发子进程。长期来看我建议垃圾分类这类小工程直接在 Linux 云服务器上跑或者用 WSL 2 也行能少踩很多环境上的坑。有一个稳定的深度学习环境比调参重要得多。5.5 推理时所有类别概率几乎相等模型像在摆烂现象用训练好的模型对单张图片做预测输出的六个概率都是 0.16 左右完全没有区分度。原因常见是两个动作漏了。一是推理时没有调用 model.eval()模型仍处于训练模式BatchNorm 还在用当前 batch 的统计信息导致输出不稳定二是对输入图片没有做相同的 Resize 和 Normalize或者直接用 OpenCV 读取的 BGR 数组传给模型通道顺序就是错的。解决推理代码里固定三件套model.eval()、torch.no_grad()、图片先走一遍 val_transform。这三件事少一件都会让结果变得像随机数。如果还是不对把输入张量打印出来看数值范围和训练时的样本是否一致。模型是个黑匣子但输入不是查输入永远是最快的排查路径。6. 把模型封装成命令行小工具做验证6.1 predict.py一个文件完成单图推理训练完模型把权重保存到 best_model.pt 之后我习惯把它封装成一个不依赖训练环境的预测脚本。后续想验证一张照片直接python predict.py photo.jpg就行不用再打开训练代码。import torch from torchvision import models, transforms from PIL import Image device torch.device(cuda if torch.cuda.is_available() else cpu) class_names [cardboard, glass, metal, paper, plastic, trash] model models.resnet18(weightsNone) model.fc torch.nn.Linear(model.fc.in_features, 6) state torch.load(best_model.pt, map_locationdevice) model.load_state_dict(state) model.to(device).eval() 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]), ]) def predict(img_path: str): img Image.open(img_path).convert(RGB) x transform(img).unsqueeze(0).to(device) with torch.no_grad(): logits model(x) prob torch.softmax(logits, dim1)[0] top2 torch.topk(prob, 2) print(预测类别:, class_names[top2.indices[0]]) for idx, conf in zip(top2.indices, top2.values): print(f{class_names[idx]}: {conf:.2%}) if top2.values[0] - top2.values[1] 0.1: print(低置信度建议人工确认后再投放) if __name__ __main__: predict(test_images/photo_01.jpg)这段代码的套路是固定的加载权重、切到 eval 模式、对输入走一遍和训练时一致的 transform。top2 的判断值得多说一句。垃圾分类的场景里用户真正需要的是“能不能扔”如果模型对前两个类别的置信度只差不到 10%说明图片本身有歧义这时候强制给出唯一答案反而危险。输出一句“建议人工确认”比硬着头皮猜更符合工程落地习惯。6.2 用手机里的真实照片做最终验收封装完脚本之后最终验收不要只在 val 集上看数字。我的做法是拿手机拍 20 张家里常见的垃圾照片覆盖不同光源、不同角度、不同遮挡程度然后统一放到 test_images 文件夹里跑一遍 predict.py。记录每张图的预测结果和人工标签算出一个“真实照片准确率”。这个数字才是你以后敢拍胸脯告诉别人的指标。我记得第一次做这个工程的时候val 集上明明有 94%自我感觉很好。后来拿手机拍了几张厨房垃圾发现一半都被判错了。那时候才意识到公开数据集里“干净”的图片和真实场景里的“脏”图片之间隔着一条巨大的沟。从那以后每次训练完都会先跑一遍真实照片测试集这个习惯帮我避免了好几回交付前的翻车。垃圾分类这个小工程的价值不在于把模型刷到多高的准确率而在于让你完整走一遍“数据整理—迁移学习—训练调参—部署验证”的流程。流程跑通了后面换成别的分类任务比如钢材缺陷、农产品分级换一下数据集和类别数就能复用。希望这些脚本和参数能帮你少踩几个坑早点跑出属于自己的那条准确率曲线。希望帮到你。本文还有配套的精品资源点击获取
返回列表