简介面向高校人工智能、计算机相关专业学生这份基于深度学习卷积神经网络的垃圾识别分类系统源码是一套可直接用于课程设计或期末大作业的完整项目。项目已获导师指导并取得97分高分涵盖从数据集准备、模型训练到分类预测的完整流程下载后无需修改即可运行。资源包共包含2000个文件其中1196张jpg与789张jpeg图像构成分类数据集13个py文件对应数据加载、模型构建、训练与推理等核心模块另有json配置与md说明文档整体压缩包约564.69MB结构清晰便于按需查阅。目前已有263人学习下载适合需要快速落地图像分类项目或参考高分结课方案的学生。借助该资源使用者可掌握CNN在垃圾分类场景中的实际应用并复用数据增强、模型保存与评估等关键代码为后续相关课题提供扎实基础。1. 垃圾识别分类系统的课程设计为什么 CNN 是默认起点如果你是做课程设计、毕业设计或者入门 CV 练手拿到“垃圾识别分类系统”这个题目第一反应应该是先把问题定义清楚输入一张图片输出一个类别——这属于典型的图像分类任务而深度学习里的卷积神经网络CNN就是这类任务的标准解法。垃圾识别分类之所以总被拿来当课程设计题目是因为它比猫狗分类多了一点真实应用背景又不像目标检测那样需要处理边框和定位数据集的敏感程度也不高适合在一两周内把“数据 → 训练 → 评估 → 展示”整条链路跑通。这套系统里Python 是粘合剂PyTorch 是主流选择数据集和预训练模型决定了你最终花多少时间在调参而不是写代码上。接下来从数据集规范讲起逐步走到模型选型、训练脚本和推理诊断。整个项目里真正的难点不是网络结构本身——CNN 的残差块、汇聚层池化层都是现成的——而是你对数据、对训练过程有没有把控力。2. 数据集准备目录规范、标签策略与数据增强2.1 数据集目录结构文件夹名就是标签拿到别人的源码包先不要急着打开训练脚本而是先看数据集长什么样。垃圾识别分类的数据集常见有两种组织方式一种是在 data/train 下按类别建子文件夹比如 cardboard / glass / metal / paper / plastic / trash另一种是给一张 CSV 表里面列了图片路径和对应标签。课程设计版本绝大多数是第一种因为 PyTorch 的torchvision.datasets.ImageFolder可以直接消费这种目录结构不需要手写 Dataset 类。我一般会先把数据目录立成这种结构data/ ├── train/ │ ├── cardboard/ │ │ ├── img_001.jpg │ │ ├── img_002.jpg │ │ └── ... │ ├── glass/ │ ├── metal/ │ ├── paper/ │ ├── plastic/ │ └── trash/ └── val/ ├── cardboard/ ├── glass/ └── ...注意一个关键点ImageFolder默认按文件夹名字母顺序排序来分配类别索引也就是说 cardboard 是 0glass 是 1依次类推。这个顺序看起来无所谓但它直接影响到你推理阶段打出来的标签对不对后期如果换数据集类别索引全部要对齐。我的习惯是加载完之后先把dataset.class_to_idx打印出来看一眼确认和我的预期一致。逻辑说明ImageFolder做的事情是遍历根目录下每个子文件夹用子文件夹名作为类别名自动生成从 0 开始的整数标签。它不需要额外的标注文件这是它适合课程设计的原因——你不用写 XML 或者 JSON 的解析代码拿到数据就能用。参数上你需要关注的就是root路径和is_valid_file如果你里面有隐藏文件或损坏图片可以用这个参数过滤。2.2 训练集与验证集的划分脚本很多下载下来的数据集只提供了一整个文件夹没有划分 train 和 val。这时候要自己写划分脚本。划分的原则是尽量保持每类图片在训练集和验证集里的比例一致最简单的方式就是按类别目录操作每个类别内部随机切分。下面这个脚本会在all_data/下每个类别文件夹里随机取 20% 的图片复制到data/val/对应目录下其余留在data/train/import os import random import shutil src_root all_data # 原始数据内部按类别分文件夹 dst_root data val_ratio 0.2 random.seed(42) # 固定随机种子保证可复现 for class_name in os.listdir(src_root): class_path os.path.join(src_root, class_name) if not os.path.isdir(class_path): continue imgs [f for f in os.listdir(class_path) if f.lower().endswith((.jpg, .jpeg, .png))] random.shuffle(imgs) val_count int(len(imgs) * val_ratio) os.makedirs(os.path.join(dst_root, train, class_name), exist_okTrue) os.makedirs(os.path.join(dst_root, val, class_name), exist_okTrue) for i, fname in enumerate(imgs): src os.path.join(class_path, fname) if i val_count: dst os.path.join(dst_root, val, class_name, fname) else: dst os.path.join(dst_root, train, class_name, fname) shutil.copy(src, dst) print(划分完成)逻辑说明按类别循环而不是全局混洗是为了防止某类图片全掉进训练集而验证集里没有这类样本。random.seed(42)的作用是让每次运行划分结果一致——课程设计答辩时如果评委问你划分细节你能明确说出一共有多少张、每类多少张这比“我随便分的”可信得多。这里用的是shutil.copy源数据不会被破坏这个习惯建议保留。参数说明val_ratio一般取 0.15 到 0.25。垃圾识别数据集通常每类几百张图片取 0.2 意味着每类留出几十张做验证够用。如果数据集特别小每类不到 100 张建议改成 0.15多留点训练样本。2.3 数据增强小数据集抗过拟合的第一道防线垃圾识别课程设计的数据量通常不会太大一个常见的现象是训练集准确率冲到 99%验证集卡在 80% 不动——这就是过拟合。对图像分类来说最直接的办法不是换模型而是做数据增强。PyTorch 里做数据增强就是往transforms里加变换注意有一个原则训练集做增强验证集只做缩放和归一化。验证集如果也做随机翻转你看到的指标就带噪声了没法稳定判断模型好坏。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(15), 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((256, 256)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])逻辑说明RandomHorizontalFlip随机水平翻转对垃圾图片是合理的——你翻转一个易拉罐它还是一个易拉罐。RandomRotation(15)是旋转 ±15 度角度别调太大超过 30 度会让“纸板”这类有文字标识的图片失真严重。ColorJitter调整亮度对比度模拟不同光照环境下拍摄的垃圾照片这是垃圾桶场景下很实际的扰动。最后Normalize用的是 ImageNet 的均值和标准差因为我们后面要用 ImageNet 预训练权重做迁移学习输入数据的分布必须和预训练时保持一致。这里有一个很多人会忽略的细节Resize((256, 256))之后网络输入分辨率是多少如果你用的是 ResNet 系列标准做法是缩放到 256 后随机裁剪 224。上面的代码里没有加RandomCrop相当于直接用 256 分辨率进网络。ResNet 的全局池化层平均汇聚层对输入尺寸不敏感所以 256 也能跑但你加载 ImageNet 预训练权重时它训练时见过的是 224 分辨率的裁剪图输入分布会对不上收敛速度会略慢。建议加上train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(15), 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((256, 256)), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])验证集用CenterCrop(224)保证每次评估时输入完全确定结果可复现。数据增强是“免费午餐”它不会增加推理成本只影响训练时每个 epoch 的数据。3. 卷积神经网络选型ResNet 与 MobileNet 的迁移学习取舍3.1 三个候选网络的对比VGG16 / ResNet18 / MobileNetV3课程设计里最常见的三个选择是 VGG16、ResNet18 和 MobileNetV3。先看它们的关键差异模型参数量224×224 CPU 单张推理耗时参考训练速度适用场景VGG16约 1.38 亿慢百毫秒级慢学习卷积和汇聚层的原理不推荐实际使用ResNet18约 1170 万中等几十毫秒较快课程设计首选均衡MobileNetV3-Small约 250 万快十毫秒级快CPU 训练或数据量很小时首选VGG16 在网络结构上最直白卷积层 汇聚层池化层堆叠适合理解 CNN 的组成方式但它的全连接层占了绝大部分参数量训练慢且容易过拟合。ResNet18 引入了残差连接解决了深层网络的梯度消失问题参数量只有 VGG16 的十分之一不到效果却更好——这也是它成为分类任务主力模型的原因。MobileNetV3 用了深度可分离卷积把计算量压缩了一个量级在纯 CPU 环境下训练也能吃得消。对于垃圾识别分类这个任务我的建议是有 NVIDIA 显卡就用 ResNet18只有 CPU 就用 MobileNetV3。不需要纠结“哪个模型准确率更高”——在课程设计的数据量下两个模型的差距很小真正的差距在数据清洗和超参数调优上。不要一上来就上 ResNet50 或 EfficientNet训练时间翻几倍准确率可能只涨一个点性价比太低。这里顺带说一句有人会问“能不能用 YOLOv8 做垃圾分类”那是目标检测路线输出的是“哪里有什么垃圾”的检测框更适合桌面级机器人抓取这种场景。纯图片分类的课程设计用分类网络就够了这也是你标题里写“卷积神经网络”而不是“目标检测”的原因——问题域决定模型选型。3.2 迁移学习代码加载预训练权重并替换分类头垃圾识别数据集和 ImageNet 完全不重叠那为什么还要用 ImageNet 预训练权重因为卷积神经网络的前几层学到的是通用特征——边缘、纹理、颜色块——这些特征对任何图像任务都有用。垃圾图片虽然和 ImageNet 里的类别不同但底层视觉结构是相似的。迁移学习的收益对小数据集尤其明显它让模型从一开始就站在一个“见过世面”的起点上而不是从随机权重开始瞎猜。PyTorch 加载预训练模型并替换分类头代码非常简洁import torchvision.models as models import torch.nn as nn num_classes 6 # 根据你的数据集类别数修改 model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) # 替换最后一层全连接层 model.fc nn.Linear(model.fc.in_features, num_classes)逻辑说明weightsmodels.ResNet18_Weights.DEFAULT会下载在 ImageNet 上预训练好的权重下载到本地缓存目录后后续每次运行都直接读缓存不会再重复下载。model.fc.in_features是 ResNet18 最后一个卷积块的输出维度恒等于 512但我们不写死 512而是从原模型里读出来——这样以后换 ResNet34512或 MobileNetV31024时这一行不用改。model.fc nn.Linear(512, num_classes)这一步是关键中的关键。如果不替换模型输出是 1000 维ImageNet 的类别数你的垃圾类别只有 6 类训练时损失函数会直接报错或者你被迫把标签扩展到 1000 维——那是灾难。参数说明weights参数在 PyTorch 1.9 之前用的写法是pretrainedTrue新版里标了废弃警告。如果你拿到别人的源码用了旧写法还能跑说明版本兼容性没问题但如果遇到加载报错优先检查这里是不是版本差异。3.3 冻结还是微调按数据集规模决定加载完预训练权重后有一个策略选择是冻结大部分卷积层只训练最后几层还是整个模型全部参与训练这决定了训练速度和收敛表现分歧很大。冻结的含义是让某些层的参数在反向传播时不更新。常见做法是for param in model.parameters(): param.requires_grad False # 只解冻最后一层卷积块和全连接层 for param in model.layer4.parameters(): param.requires_grad True for param in model.fc.parameters(): param.requires_grad True # 优化器只传入 requires_gradTrue 的参数 optimizer torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr1e-3 )逻辑说明requires_gradFalse表示这个参数不计算梯度也不参与更新。冻结前面所有层、只训练layer4和fc相当于只训练整个网络最后 20% 的部分。它的优点是训练速度快、显存占用低适合每类只有几十张图片的极端小数据集缺点是因为前面的特征提取器不动模型对垃圾图片特有的特征比如塑料的透明反光适应能力弱。我的经验是分两步走如果数据集每类图片不足 100 张先全部冻结只训fc层跑 10 个 epoch看看验证集准确率能到多少如果准确率已经能接受就不折腾了。如果准确率不够再解冻layer4甚至全部层用一个更小的学习率1e-4微调。这个策略能在不显著增加训练时间的前提下把准确率抬上去。课程设计答辩时能说清楚“为什么先冻结后微调”比直接整网训练更能体现对模型的理解。4. 训练启动从超参数到训练脚本的落地配置4.1 最小可用的训练脚本网络结构定了、数据加载定了剩下的就是把训练循环跑起来。这里给一个可以直接改用的训练脚本骨架它包含了训练、验证、保存最优模型和动态调整学习率这几个核心动作。import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, transforms import copy device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) train_dataset datasets.ImageFolder(data/train, transformtrain_transform) val_dataset datasets.ImageFolder(data/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30) best_acc 0.0 best_state None for epoch in range(30): # 训练阶段 model.train() running_loss 0.0 for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) # 验证阶段 model.eval() correct 0 total 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) acc correct / total print(fEpoch {epoch1:02d} | Loss {running_loss/len(train_dataset):.4f} | Val Acc {acc:.4f}) if acc best_acc: best_acc acc best_state copy.deepcopy(model.state_dict()) scheduler.step() torch.save(best_state, best_model.pth)逻辑说明model.train()和model.eval()切换的是 BatchNorm 和 Dropout 的行为——训练时 BN 用当前 batch 的均值方差评估时用训练阶段累计的全局统计量。忘记切eval()是验证集准确率忽高忽低的常见原因BatchNorm 的坑比别处多。验证阶段套了torch.no_grad()因为验证不需要计算梯度不套它会多占用一倍显存。copy.deepcopy保存的是最佳权重而不是引用防止后续 epoch 修改了同一份内存导致“最佳模型”名存实亡。参数说明batch_size32是 1080Ti 及以上显存11GB的安全值如果显卡只有 4GB 就降成 16CPU 训练建议 8。num_workers4是数据加载的并行进程数Windows 下如果报错改成 0。T_max30要和总 epoch 数一致它表示余弦退火学习率的一个周期长度——0 到 30 个 epoch学习率从初始值平滑降到接近 0。4.2 超参数配置表与选择理由训练神经网络最怕“参数全默认”因为默认值是为一般场景设计的不是你数据集的。下面这张表是我做垃圾分类训练时常用的配置直接抄能用超参数推荐值理由输入分辨率224×224与 ImageNet 预训练输入一致epochs30课程设计数据量下30 轮足够收敛batch_size32GPU/ 8CPU太小噪声大太大显存放不下优化器AdamWAdam 加了权重衰减解耦收敛稳定初始学习率1e-4微调/ 1e-3重新训练迁移学习用大学习率会把预训练权重冲掉weight_decay1e-4轻微 L2 正则抑制过拟合学习率调度CosineAnnealingLR前期快速下降后期精细收敛比固定学习率好损失函数CrossEntropyLoss多分类标准配置这里重点说两个。第一个是学习率用 ImageNet 预训练权重做迁移学习初始学习率不要超过 1e-3我见过太多人用 1e-2 起步loss 直接变成 NaN 或者震荡不停最后归因于“深度学习是玄学”——其实只是学习率不适配。第二个是weight_decay它不是可有可无的装饰品垃圾数据集小过拟合风险本来就高1e-4 的权重衰减相当于给参数套了个约束能有效控制训练集和验证集准确率的差距。4.3 训练过程监控loss 不降和过拟合怎么判断训练不是把脚本扔进终端就不管了。前 5 个 epoch 是最关键的观察窗口大部分问题这时候就会暴露。如果 loss 根本没下降或者从第一个 epoch 开始就是 NaN按顺序排查两件事一是学习率是否过大调小 10 倍试试二是数据里是否有损坏图片PIL 解码失败返回 NoneNormalize 过不了。如果 loss 在下降但验证集准确率震荡剧烈常见原因是 batch_size 太小导致梯度噪声过大或者验证集图片太少每类不到 20 张准确率波动就是假的波动。如果训练集准确率一路冲高到 98%验证集却卡在 75% 上下那就是典型的过拟合。这时候不要着急换模型先做三件事加强数据增强把 RandomRotation 从 15 加到 20加上 RandomResizedCrop、增大 weight_decay 到 1e-3、早停训练集准确率持续高于验证集 15 个百分点以上时就停止保存之前的最佳模型。这三个动作处理掉大部分过拟合场景。5. 垃圾分类训练与推理中常见的 5 个翻车点排查5.1 中文路径与带空格文件夹名导致读图失败现象用 OpenCV 的cv2.imread读取图片返回的数组是None或者ImageFolder能加载目录但训练时报错说找不到文件。原因OpenCV 的imread底层用的是 C 的fopen不处理 Unicode 编码Windows 下遇到中文路径或中文文件名直接返回空。而课程设计下载的数据集很多就是从百度图片里批量爬的文件名带中文标签非常常见。解决不要用cv2.imread改用 PIL 读取或者用 OpenCV 的“绕路”方法import cv2 import numpy as np def imread_unicode(path): data np.fromfile(path, dtypenp.uint8) return cv2.imdecode(data, cv2.IMREAD_COLOR)np.fromfile以二进制方式读取文件不受路径编码限制再交给cv2.imdecode解码成图像。更省事的方案是在数据集预处理阶段把所有文件名批量改成纯英文数字img_0001.jpg这种格式一劳永逸地消灭编码问题。我的建议是别在训练代码里做兼容直接在数据清洗阶段把文件名改干净后面所有代码都不用再为编码操心。5.2 灰度图和 RGBA 图混入数据集训练直接崩溃现象训练跑到一半报错信息类似 “expected 3 input channels, got 1”或者 “不能将 size 为 4 的维度广播为 3”。原因垃圾识别数据集经常是多个渠道收集拼凑的有些是手机拍的3 通道 RGB有些可能是扫描件或老照片灰度图 1 通道还有透明背景的 PNG4 通道 RGBA。这三种图堆在一个文件夹里第一个 batch 可能没事第二个 batch 遇到灰度图就崩。解决在自定义 Dataset 的__getitem__里强制统一通道数或者更简单——在数据增强前加一步from PIL import Image def process_image(path): img Image.open(path).convert(RGB) return imgconvert(RGB)会把灰度图复制成三通道把 RGBA 图丢弃透明通道。这是最省事的方案不需要写额外的判断逻辑。注意要在一个地方统一处理不要一个脚本里用 PIL、另一个脚本里用 OpenCV两边的通道顺序都不一致PIL 是 RGBOpenCV 是 BGR混用就会得到颜色完全错乱的增强图。5.3 预训练权重下载卡住训练一直停在开始现象第一次运行训练脚本程序卡在 “Downloading” 字样进度条一直不动或者下载到一半报连接超时。原因torchvision 的预训练权重托管在外部服务器上网络环境不稳定时下载大文件ResNet18 约 45MB容易中断。而且 PyTorch 的下载逻辑不会断点续传中断后留下一个.tmp后缀的临时文件下次运行以为下载过了实际是个损坏文件。解决先把临时文件清干净再手动下载权重放到本地缓存目录。缓存位置一般在~/.cache/torch/hub/checkpoints/Windows 在C:\Users\用户名\.cache\torch\hub\checkpoints\。操作步骤# 清理残留的临时文件 rm -f ~/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth.tmp # 手动下载权重文件到该目录后重命名为完整文件名之后在代码里加载ResNet18_Weights.DEFAULT时程序会发现缓存文件已存在直接跳过下载。这个排查方法同样适用于加载其他预训练模型。如果下载总是失败还有一个不依赖网络的备选思路自己从头训练把weightsNone但准确率会明显低几个百分点而且需要更长训练时间不推荐作为首选方案。5.4 loss 降了但验证集准确率纹丝不动现象训练 loss 每轮都在降看起来一切正常但验证集准确率从第 5 轮开始就卡在一个位置怎么训都不涨。最典型的卡点是 50% 左右——二分类或者类别极度不均衡时就是“永远猜多数类”的准确率。原因最常见的是类别不均衡。垃圾分类数据集如果按“可回收 / 厨余 / 有害 / 其他”四分类有害垃圾的图片可能只有几十张其他三类各有几百张。模型学到的是偏向多数类的决策边界loss 还在降只是多数类拟合得更好而已。另一个常见原因是验证集划分不合理比如同源图片被同时分进了训练集和验证集——同一物品的连拍图片。解决训练前先打印每个类别的图片数量统计。如果发现某类图片特别少有三个方案一是用WeightedRandomSampler让每个 batch 里少数类图片出现的概率更高二是把少数类做更强的数据增强比如专门对少数类做 90 度、180 度旋转三是最简单粗暴的——增加该类图片。课程设计阶段建议用第三个方案多拍几十张照片不费事效果比调采样器靠谱。from torch.utils.data import WeightedRandomSampler labels [s for _, s in train_dataset.samples] class_counts torch.bincount(torch.tensor(labels)) weights 1.0 / class_counts.float() sample_weights weights[labels] sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights), replacementTrue) train_loader DataLoader(train_dataset, batch_size32, samplersampler)加权采样的原理每个样本被抽中的概率和它的sample_weights成正比少数类的1 / count更大所以被采到的次数更多。这个方案效果好但要注意sampler和shuffleTrue不能同时用DataLoader里要删掉shuffle参数。5.5 换数据集后忘了同步改分类头验证集准确率全零现象代码之前跑通过一个 6 类垃圾分类数据集后来换了一个 40 类的“华为云垃圾分类”数据集重跑训练训练不会报错但验证集准确率一直是 0或者推理时输出的预测结果全是同一个类别。原因换了数据集以后num_classes没从 6 改成 40model.fc还是 6 维输出。训练脚本里torch.max(outputs, 1)取的是 6 个类别里的最大概率而标签范围是 0 到 39绝大概率标签不在预测范围内准确率当然接近 0。这种情况训练不报错因为CrossEntropyLoss对标签的上界没有运行时检查——它只会算出一个看起来还合理的 loss。解决每次换数据集后做两件事。第一确认num_classes和ImageFolder自动生成的class_to_idx长度一致num_classes len(train_dataset.classes) model.fc nn.Linear(model.fc.in_features, num_classes)第二训练完成后打印一张类别对照表保存下来这样推理时做标签到类名的映射就不会错位。代码里如果之前写了category_names [cardboard, glass, ...]这种硬编码列表换数据集后必须同步更新——这是课程设计源码里最隐蔽的坑之一表面上看是“模型不行”实际是标签映射断了。6. 推理误判诊断用混淆矩阵和类别召回率调模型6.1 单张图片推理的最小代码训练完模型最后一步是写推理脚本。加载权重时注意保存的是state_dict纯权重还是整个模型权重加结构两者加载方式不同。之前的代码保存的是state_dict推理脚本里需要先创建模型实例再往里灌权重。import torch from PIL import Image from torchvision import transforms device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet18(weightsNone) model.fc torch.nn.Linear(512, num_classes) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model.to(device).eval() class_names train_dataset.classes # 训练时保存下来的类别名列表 def predict(image_path): img Image.open(image_path).convert(RGB) img_tensor val_transform(img).unsqueeze(0).to(device) with torch.no_grad(): outputs model(img_tensor) prob torch.softmax(outputs, dim1) pred_idx torch.argmax(prob, dim1).item() return class_names[pred_idx], prob[0][pred_idx].item() print(predict(test/glass_001.jpg))逻辑说明map_locationdevice是必须的——如果模型在 GPU 上训练保存的权重张量在 CUDA 设备上纯 CPU 环境加载时没有map_locationcpu就会报错。推理时也要走val_transform尤其是CenterCrop和Normalize不能为了省事只做Resize否则图片缩放和归一化的分布和训练时不一致预测结果会偏。6.2 误分类网格图定位是数据问题还是模型问题训练完不要只看一个总体准确率就收工尤其是答辩时老师大概率会问“什么类型的垃圾容易分错”。我的习惯是用一个可视化脚本把验证集里预测错误的图片按类别分组输出成一张网格图肉眼判断错误模式。import matplotlib.pyplot as plt def show_misclassified(model, val_loader, class_names, max_show16): mis [] with torch.no_grad(): for inputs, labels in val_loader: outputs model(inputs) preds torch.argmax(outputs, dim1) for i in range(len(labels)): if preds[i] ! labels[i]: mis.append((inputs[i], labels[i].item(), preds[i].item())) if len(mis) max_show: break fig, axes plt.subplots(4, 4, figsize(12, 12)) for idx, (img, true, pred) in enumerate(mis[:16]): img img.permute(1, 2, 0).numpy() img img * [0.229, 0.224, 0.225] [0.485, 0.456, 0.406] img img.clip(0, 1) axes[idx // 4][idx % 4].imshow(img) axes[idx // 4][idx % 4].set_title(fTrue: {class_names[true]}\nPred: {class_names[pred]}) axes[idx // 4][idx % 4].axis(off) plt.tight_layout() plt.savefig(misclassified.png, dpi150)如果网格图里显示的错误集中在某两对类别上比如玻璃 vs 塑料说明这两个类别的视觉特征确实相近这时可以去看混淆矩阵计算每类的召回率——如果“玻璃”的召回率只有 60%而其他类都在 90% 以上就值得单开一组数据增强策略或者多补拍玻璃类的图片。这是我能给的最实用的调模型习惯不要调模型结构先看数据哪里不足。我自己的习惯是每次训练任务跑完先花 10 分钟看错误网格图再决定要不要继续加 epoch 或者调学习率。很多次我以为“模型有问题”最后发现是数据里混了标注错误的图片——标签错了再怎么调参都白搭。把这张“错误路径图”当成你调参的起点而不是把准确率数字当成终点能少走很多弯路。希望帮到你。本文还有配套的精品资源点击获取
