简介这份资源面向计算机类毕业设计与课程作业场景聚焦深度学习中的垃圾分类实战任务适合具备一定Python基础、正在寻找完整项目参考的本科生或自学者。包内共6个文件以2个py源码、2个pt权重文件为主另含1个zip压缩包与1份md说明文档整体约4MB体量轻便便于快速部署与二次修改。源码部分涵盖模型定义与预测推理脚本权重文件可直接加载用于验证CNN与ResNet两种网络结构的分类效果md文档则提供项目说明与使用指引。已有77人学习下载说明该方案在同类毕设选题中具备一定参考价值。读者可借此获得一套可运行的垃圾分类识别流程理解卷积神经网络与残差网络在图像分类任务中的差异并在此基础上调整网络结构、替换数据集或补充训练脚本完成从模型搭建到推理验证的完整闭环为毕设答辩或课程提交提供扎实的代码与实验支撑。1. 垃圾分类深度学习毕设包从 CNN 到 ResNet 的完整落地路径很多同学做毕设时卡在同一个地方算法原理背得滚瓜烂熟真到要交一个能跑起来、能演示、能写进论文的系统时却不知道从哪下手。这个压缩包解决的正是这个问题——它把垃圾分类这个经典课题的完整链路打包好了包含model.py、model1.pt、model.pt、predict.py和README.md覆盖了从模型定义到推理预测的全流程。垃圾分类本身是深度学习图像识别的入门级任务但它的工程结构和大项目完全一致数据预处理、模型搭建、训练保存、加载推理。适合正在做深度学习毕设或课程作业的本科生也适合想找一个干净模板快速验证自己想法的开发者。包里同时提供了 CNN 和 ResNet 两条技术路线你可以根据算力和精度要求灵活选择。2. 拆开压缩包文件结构与两条模型路线的选型逻辑拿到一个源码包第一件事不是急着跑predict.py而是先搞清楚每个文件在系统里扮演什么角色。这个包的结构很典型没有多余的目录嵌套说明作者是按“最小可运行系统”的思路组织的。2.1 文件清单与各自职责文件类型作用model.py源码定义 CNN 和 ResNet 两种网络结构model.pt权重训练好的模型参数文件主模型model1.pt权重另一个版本的模型参数备用/对比predict.py源码加载权重对单张图片做推理README.md文档环境依赖和基本使用说明Graduation Design validation.zip压缩包验证集数据或相关验证脚本model.py是整个项目的核心它决定了网络长什么样。model.pt和model1.pt是两个已经训练好的权重文件意味着你不需要从零训练就能直接跑推理。predict.py是入口脚本负责把图片喂给模型并输出分类结果。这种“定义-权重-推理”三件套的结构在深度学习项目里是最常见的组织方式。2.2 CNN 和 ResNet 该怎么选包里的model.py同时提供了 CNN 和 ResNet 两种结构这不是凑数而是对应了两种不同的使用场景。CNN 版本通常是几层卷积加池化的堆叠参数量小训练快适合以下情况你的电脑没有独立显卡只能用 CPU 跑或者你需要在论文里展示一个“从零搭建”的网络结构体现你对卷积核、池化、全连接层的理解。CNN 的缺点是层数一深就容易出现梯度消失精度上限有限。ResNet 版本引入了残差连接核心思想是让网络学习输入和输出之间的“残差”而不是直接学习映射。这样做的好处是梯度可以通过跳跃连接直接回传使得训练几十层甚至上百层的网络成为可能。垃圾分类这种任务类别之间的区分度主要在颜色、纹理和形状上ResNet 的深层特征提取能力明显更强。如果你有 GPU 可用或者追求更高的分类准确率直接用 ResNet 版本。提示如果你不确定选哪个先跑predict.py看默认加载的是哪个权重文件再决定要不要切换。2.3 环境依赖与快速验证在跑任何代码之前先把环境配好。这个项目依赖 PyTorch 和 torchvision另外需要 Pillow 做图像读取。常见做法是用 conda 建一个独立环境避免和系统里的其他包冲突。# 创建并激活虚拟环境 conda create -n garbage_cls python3.9 -y conda activate garbage_cls # 安装核心依赖 pip install torch torchvision pillow numpy安装完成后进入项目目录直接运行推理脚本验证环境是否正常python predict.py --image test.jpg如果predict.py没有提供命令行参数接口那就需要打开脚本找到图片路径的变量手动改成你要测试的图片路径。这一步的目的是确认模型能加载、图片能读取、输出能打印。跑通了再往下做定制化修改。3. 模型定义与推理脚本从 model.py 到 predict.py 的代码拆解环境跑通只是第一步真正要拿这个包做毕设你得能看懂model.py里网络是怎么搭的以及predict.py是怎么把图片变成分类结果的。这一章把这两个文件的核心逻辑拆开讲。3.1 model.py 里的 CNN 结构CNN 版本的典型写法是继承nn.Module在__init__里定义层在forward里定义数据流向。下面是一个符合这个项目场景的 CNN 结构示例import torch.nn as nn import torch.nn.functional as F class SimpleCNN(nn.Module): def __init__(self, num_classes4): super(SimpleCNN, self).__init__() # 第一层卷积输入3通道RGB输出32通道特征图卷积核3x3 self.conv1 nn.Conv2d(3, 32, kernel_size3, padding1) # 第二层卷积32通道扩展到64通道 self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) # 第三层卷积64通道扩展到128通道 self.conv3 nn.Conv2d(64, 128, kernel_size3, padding1) # 最大池化2x2窗口步长2特征图尺寸减半 self.pool nn.MaxPool2d(2, 2) # 全连接层128*28*28是经过三次池化后的特征维度 self.fc1 nn.Linear(128 * 28 * 28, 256) self.fc2 nn.Linear(256, num_classes) self.dropout nn.Dropout(0.5) def forward(self, x): # 每层卷积后接ReLU激活再池化 x self.pool(F.relu(self.conv1(x))) # 224-112 x self.pool(F.relu(self.conv2(x))) # 112-56 x self.pool(F.relu(self.conv3(x))) # 56-28 x x.view(-1, 128 * 28 * 28) # 展平 x F.relu(self.fc1(x)) x self.dropout(x) # 防止过拟合 x self.fc2(x) return x这段代码的逻辑很直白三层卷积逐级提取特征每层后面跟池化把特征图尺寸缩小最后展平送进全连接层做分类。num_classes4对应垃圾分类的四个类别通常是可回收、厨余、有害、其他。padding1保证卷积后特征图尺寸不变只有池化才会降维。Dropout(0.5)是训练时的正则化手段推理时会自动关闭。参数调整的关键点如果你用的图片尺寸不是 224×224全连接层的输入维度需要重新计算。计算方法是输入尺寸经过三次池化每次除以2得到最终特征图边长再乘以通道数。比如输入 128×128三次池化后是 16×16全连接层就是128 * 16 * 16。3.2 ResNet 版本的加载方式ResNet 版本通常不从头定义网络而是用 torchvision 提供的预训练模型做迁移学习。这样做的好处是预训练权重已经在大规模数据集上学到了通用特征你只需要替换最后的全连接层用垃圾分类数据微调即可。import torchvision.models as models import torch.nn as nn def get_resnet_model(num_classes4): # 加载ResNet18预训练权重 model models.resnet18(pretrainedTrue) # 冻结前面的层只训练最后的分类层 for param in model.parameters(): param.requires_grad False # 替换全连接层ResNet18的fc输入是512维 model.fc nn.Linear(512, num_classes) return modelpretrainedTrue表示加载在 ImageNet 上训练好的权重。requires_grad False冻结了卷积层的参数训练时只更新fc层。如果你的数据集比较小几千张以内这种做法能有效防止过拟合。如果数据量够大可以把requires_grad设为True做全网络微调但学习率要调小。3.3 predict.py 的推理流程推理脚本的核心步骤是加载模型结构 → 加载权重 → 图片预处理 → 前向传播 → 输出类别。下面是一个标准的推理流程import torch from PIL import Image from torchvision import transforms from model import SimpleCNN # 或 get_resnet_model # 定义预处理和训练时保持一致 transform transforms.Compose([ transforms.Resize((224, 224)), # 统一尺寸 transforms.ToTensor(), # 转成Tensor归一化到[0,1] transforms.Normalize( # 标准化均值方差用ImageNet的 mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) ]) # 加载模型 model SimpleCNN(num_classes4) model.load_state_dict(torch.load(model.pt, map_locationcpu)) model.eval() # 切换到推理模式关闭Dropout和BatchNorm的更新 # 读取并预处理图片 image Image.open(test.jpg).convert(RGB) input_tensor transform(image).unsqueeze(0) # 增加batch维度 # 推理 with torch.no_grad(): # 关闭梯度计算节省内存 output model(input_tensor) _, predicted torch.max(output, 1) classes [可回收, 厨余, 有害, 其他] print(f预测结果{classes[predicted.item()]})model.eval()和torch.no_grad()是推理时必加的两行。前者关闭 Dropout 和 BatchNorm 的训练行为后者关闭梯度计算。unsqueeze(0)是因为模型期望的输入是[batch, channel, height, width]单张图片只有后三个维度需要手动补一个 batch 维度。map_locationcpu保证在没有 GPU 的机器上也能加载权重。注意预处理中的Normalize参数必须和训练时一致。如果训练时没用标准化推理时也不要加否则输入分布不匹配预测结果会完全错误。4. 避坑与排查权重加载、类别映射和预处理不一致的常见翻车这个包虽然结构简单但实际跑起来翻车的概率不低。下面是我在帮别人看毕设代码时最常遇到的几个问题按“现象 → 原因 → 解决”整理。4.1 权重加载报错Missing key(s) in state_dict现象运行predict.py时抛出RuntimeError: Error(s) in loading state_dict for SimpleCNN: Missing key(s)...。原因model.py里的网络结构和model.pt保存时的结构不一致。常见情况是作者训练时用了 ResNet但推理时实例化的是 CNN或者改了全连接层的输出维度。解决先打印权重文件的键名和当前模型对比state_dict torch.load(model.pt, map_locationcpu) for k, v in state_dict.items(): print(k, v.shape)对比model.py里定义的层名找到不匹配的地方。如果是全连接层维度不对修改num_classes或fc层的输入维度。如果是层名对不上说明模型结构定义有差异需要根据权重文件的键名反推网络结构。4.2 预测结果全是同一类现象不管输入什么图片输出的类别都一样。原因预处理不一致是最常见的原因。训练时用了Normalize推理时没加或者训练时图片是 BGR 格式推理时用了 RGB。另一个可能是权重文件本身就没训练好。解决先检查predict.py里的transform是否和训练脚本一致。如果找不到训练脚本可以尝试去掉Normalize再跑一次看结果是否变化。如果还是全同一类用torch.load检查权重里的数值是否正常有没有大量零值或 NaN。4.3 CUDA out of memory现象在有 GPU 的机器上跑推理时报显存不足。原因predict.py默认把模型和输入都放到了 GPU 上但显存被其他进程占用了或者图片尺寸太大。解决在加载权重时加map_locationcpu强制用 CPU 推理。或者在推理前加torch.cuda.empty_cache()清理缓存。如果必须用 GPU把输入图片的 batch size 设为 1不要一次加载多张。4.4 类别标签对不上现象模型输出0但classes列表里0对应的类别和实际不符。原因训练时类别的索引顺序和推理时classes列表的顺序不一致。比如训练时0是可回收推理时写成了0是厨余。解决找到训练时的类别映射关系。通常在训练脚本里会有一个class_to_idx的字典或者在数据加载时用ImageFolder会自动按文件夹名排序生成索引。以那个顺序为准不要凭记忆写。4.5 model.pt 和 model1.pt 搞混现象加载model.pt效果正常加载model1.pt效果很差。原因两个权重文件对应不同的模型结构或不同的训练轮数。model1.pt可能是早期保存的检查点或者对应的是 CNN 而不是 ResNet。解决分别用两个权重跑同一张测试图对比输出。如果model1.pt效果差检查它对应的模型结构是否和当前实例化的模型一致。在README.md里通常会说明哪个是最终版本。5. 从能跑到能用替换数据集、调整类别数和验证模型效果的实操技巧把预训练权重跑通只是验证了环境真正要做自己的毕设你得把模型适配到自己的数据集上。这一章讲怎么替换数据、调整类别数以及怎么验证模型是不是真的学到了东西。5.1 用自己的数据重新训练假设你的数据集按文件夹组织每个类别一个文件夹dataset/ ├── recyclable/ ├── kitchen_waste/ ├── hazardous/ └── other/用ImageFolder加载数据然后微调 ResNetfrom torchvision import datasets, transforms from torch.utils.data import DataLoader # 数据增强训练时随机裁剪、翻转增加泛化能力 train_transform transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(dataset/, transformtrain_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue) # 类别数自动从数据集获取 num_classes len(train_dataset.classes) print(f类别{train_dataset.classes}) # 打印索引映射关系ImageFolder会按文件夹名的字母顺序生成索引这个顺序就是推理时classes列表应该用的顺序。batch_size32是常见起点显存不够就降到 16 或 8。shuffleTrue打乱训练顺序防止模型学到数据排列的规律。5.2 训练循环的关键参数import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model get_resnet_model(num_classes).to(device) criterion nn.CrossEntropyLoss() # 只优化fc层的参数学习率设大一点 optimizer optim.Adam(model.fc.parameters(), lr0.001) for epoch in range(10): 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() print(fEpoch {epoch1}, Loss: {running_loss/len(train_loader):.4f}) # 保存权重 torch.save(model.state_dict(), my_model.pt)CrossEntropyLoss是分类任务的标准损失函数内部包含了 softmax。Adam优化器对学习率不敏感lr0.001是安全起点。如果只训练fc层10 个 epoch 通常就能收敛。如果做全网络微调学习率要降到0.0001左右否则预训练权重会被破坏。5.3 验证模型效果的三个检查点训练完不能只看 loss 下降就完事得做实际验证。第一看混淆矩阵。用验证集跑一遍统计每个类别的正确率和误判方向。如果“厨余”大量被误判为“其他”说明这两个类别的特征区分度不够需要增加数据或调整网络。第二看单张图片的置信度。torch.max返回的不仅是类别索引还有置信度分数。如果所有图片的置信度都在 0.3 左右说明模型没学到东西大概率是预处理或标签映射出了问题。第三用训练时没见过的图片测试。从网上找几张垃圾分类的图片手动跑predict.py看输出是否合理。这一步能暴露过拟合问题——训练集准确率很高但新图片全错。5.4 类别数调整的注意事项如果你把四分类改成二分类比如只分可回收和其他需要改两个地方model.py里的num_classes参数以及推理时的classes列表。如果用的是 ResNetmodel.fc nn.Linear(512, num_classes)会自动适配。但要注意如果加载的是原来四分类的权重全连接层的形状不匹配需要重新训练或者只加载卷积层的权重。# 只加载卷积层权重跳过fc层 pretrained_dict torch.load(model.pt) model_dict model.state_dict() # 过滤掉形状不匹配的键 pretrained_dict {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape model_dict[k].shape} model_dict.update(pretrained_dict) model.load_state_dict(model_dict)这段代码的作用是把预训练权重里和当前模型形状匹配的层加载进来不匹配的层比如fc跳过用随机初始化。这样可以在小数据集上快速启动训练。从那以后我每次拿到一个深度学习源码包都强制先跑一遍predict.py确认推理链路通畅再动model.py改结构。顺序反了的话出了问题你分不清是权重的问题还是代码改错了。希望帮到你。本文还有配套的精品资源点击获取
垃圾分类深度学习毕设实战：CNN与ResNet模型落地与避坑指南
简介这份资源面向计算机类毕业设计与课程作业场景聚焦深度学习中的垃圾分类实战任务适合具备一定Python基础、正在寻找完整项目参考的本科生或自学者。包内共6个文件以2个py源码、2个pt权重文件为主另含1个zip压缩包与1份md说明文档整体约4MB体量轻便便于快速部署与二次修改。源码部分涵盖模型定义与预测推理脚本权重文件可直接加载用于验证CNN与ResNet两种网络结构的分类效果md文档则提供项目说明与使用指引。已有77人学习下载说明该方案在同类毕设选题中具备一定参考价值。读者可借此获得一套可运行的垃圾分类识别流程理解卷积神经网络与残差网络在图像分类任务中的差异并在此基础上调整网络结构、替换数据集或补充训练脚本完成从模型搭建到推理验证的完整闭环为毕设答辩或课程提交提供扎实的代码与实验支撑。1. 垃圾分类深度学习毕设包从 CNN 到 ResNet 的完整落地路径很多同学做毕设时卡在同一个地方算法原理背得滚瓜烂熟真到要交一个能跑起来、能演示、能写进论文的系统时却不知道从哪下手。这个压缩包解决的正是这个问题——它把垃圾分类这个经典课题的完整链路打包好了包含model.py、model1.pt、model.pt、predict.py和README.md覆盖了从模型定义到推理预测的全流程。垃圾分类本身是深度学习图像识别的入门级任务但它的工程结构和大项目完全一致数据预处理、模型搭建、训练保存、加载推理。适合正在做深度学习毕设或课程作业的本科生也适合想找一个干净模板快速验证自己想法的开发者。包里同时提供了 CNN 和 ResNet 两条技术路线你可以根据算力和精度要求灵活选择。2. 拆开压缩包文件结构与两条模型路线的选型逻辑拿到一个源码包第一件事不是急着跑predict.py而是先搞清楚每个文件在系统里扮演什么角色。这个包的结构很典型没有多余的目录嵌套说明作者是按“最小可运行系统”的思路组织的。2.1 文件清单与各自职责文件类型作用model.py源码定义 CNN 和 ResNet 两种网络结构model.pt权重训练好的模型参数文件主模型model1.pt权重另一个版本的模型参数备用/对比predict.py源码加载权重对单张图片做推理README.md文档环境依赖和基本使用说明Graduation Design validation.zip压缩包验证集数据或相关验证脚本model.py是整个项目的核心它决定了网络长什么样。model.pt和model1.pt是两个已经训练好的权重文件意味着你不需要从零训练就能直接跑推理。predict.py是入口脚本负责把图片喂给模型并输出分类结果。这种“定义-权重-推理”三件套的结构在深度学习项目里是最常见的组织方式。2.2 CNN 和 ResNet 该怎么选包里的model.py同时提供了 CNN 和 ResNet 两种结构这不是凑数而是对应了两种不同的使用场景。CNN 版本通常是几层卷积加池化的堆叠参数量小训练快适合以下情况你的电脑没有独立显卡只能用 CPU 跑或者你需要在论文里展示一个“从零搭建”的网络结构体现你对卷积核、池化、全连接层的理解。CNN 的缺点是层数一深就容易出现梯度消失精度上限有限。ResNet 版本引入了残差连接核心思想是让网络学习输入和输出之间的“残差”而不是直接学习映射。这样做的好处是梯度可以通过跳跃连接直接回传使得训练几十层甚至上百层的网络成为可能。垃圾分类这种任务类别之间的区分度主要在颜色、纹理和形状上ResNet 的深层特征提取能力明显更强。如果你有 GPU 可用或者追求更高的分类准确率直接用 ResNet 版本。提示如果你不确定选哪个先跑predict.py看默认加载的是哪个权重文件再决定要不要切换。2.3 环境依赖与快速验证在跑任何代码之前先把环境配好。这个项目依赖 PyTorch 和 torchvision另外需要 Pillow 做图像读取。常见做法是用 conda 建一个独立环境避免和系统里的其他包冲突。# 创建并激活虚拟环境 conda create -n garbage_cls python3.9 -y conda activate garbage_cls # 安装核心依赖 pip install torch torchvision pillow numpy安装完成后进入项目目录直接运行推理脚本验证环境是否正常python predict.py --image test.jpg如果predict.py没有提供命令行参数接口那就需要打开脚本找到图片路径的变量手动改成你要测试的图片路径。这一步的目的是确认模型能加载、图片能读取、输出能打印。跑通了再往下做定制化修改。3. 模型定义与推理脚本从 model.py 到 predict.py 的代码拆解环境跑通只是第一步真正要拿这个包做毕设你得能看懂model.py里网络是怎么搭的以及predict.py是怎么把图片变成分类结果的。这一章把这两个文件的核心逻辑拆开讲。3.1 model.py 里的 CNN 结构CNN 版本的典型写法是继承nn.Module在__init__里定义层在forward里定义数据流向。下面是一个符合这个项目场景的 CNN 结构示例import torch.nn as nn import torch.nn.functional as F class SimpleCNN(nn.Module): def __init__(self, num_classes4): super(SimpleCNN, self).__init__() # 第一层卷积输入3通道RGB输出32通道特征图卷积核3x3 self.conv1 nn.Conv2d(3, 32, kernel_size3, padding1) # 第二层卷积32通道扩展到64通道 self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) # 第三层卷积64通道扩展到128通道 self.conv3 nn.Conv2d(64, 128, kernel_size3, padding1) # 最大池化2x2窗口步长2特征图尺寸减半 self.pool nn.MaxPool2d(2, 2) # 全连接层128*28*28是经过三次池化后的特征维度 self.fc1 nn.Linear(128 * 28 * 28, 256) self.fc2 nn.Linear(256, num_classes) self.dropout nn.Dropout(0.5) def forward(self, x): # 每层卷积后接ReLU激活再池化 x self.pool(F.relu(self.conv1(x))) # 224-112 x self.pool(F.relu(self.conv2(x))) # 112-56 x self.pool(F.relu(self.conv3(x))) # 56-28 x x.view(-1, 128 * 28 * 28) # 展平 x F.relu(self.fc1(x)) x self.dropout(x) # 防止过拟合 x self.fc2(x) return x这段代码的逻辑很直白三层卷积逐级提取特征每层后面跟池化把特征图尺寸缩小最后展平送进全连接层做分类。num_classes4对应垃圾分类的四个类别通常是可回收、厨余、有害、其他。padding1保证卷积后特征图尺寸不变只有池化才会降维。Dropout(0.5)是训练时的正则化手段推理时会自动关闭。参数调整的关键点如果你用的图片尺寸不是 224×224全连接层的输入维度需要重新计算。计算方法是输入尺寸经过三次池化每次除以2得到最终特征图边长再乘以通道数。比如输入 128×128三次池化后是 16×16全连接层就是128 * 16 * 16。3.2 ResNet 版本的加载方式ResNet 版本通常不从头定义网络而是用 torchvision 提供的预训练模型做迁移学习。这样做的好处是预训练权重已经在大规模数据集上学到了通用特征你只需要替换最后的全连接层用垃圾分类数据微调即可。import torchvision.models as models import torch.nn as nn def get_resnet_model(num_classes4): # 加载ResNet18预训练权重 model models.resnet18(pretrainedTrue) # 冻结前面的层只训练最后的分类层 for param in model.parameters(): param.requires_grad False # 替换全连接层ResNet18的fc输入是512维 model.fc nn.Linear(512, num_classes) return modelpretrainedTrue表示加载在 ImageNet 上训练好的权重。requires_grad False冻结了卷积层的参数训练时只更新fc层。如果你的数据集比较小几千张以内这种做法能有效防止过拟合。如果数据量够大可以把requires_grad设为True做全网络微调但学习率要调小。3.3 predict.py 的推理流程推理脚本的核心步骤是加载模型结构 → 加载权重 → 图片预处理 → 前向传播 → 输出类别。下面是一个标准的推理流程import torch from PIL import Image from torchvision import transforms from model import SimpleCNN # 或 get_resnet_model # 定义预处理和训练时保持一致 transform transforms.Compose([ transforms.Resize((224, 224)), # 统一尺寸 transforms.ToTensor(), # 转成Tensor归一化到[0,1] transforms.Normalize( # 标准化均值方差用ImageNet的 mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) ]) # 加载模型 model SimpleCNN(num_classes4) model.load_state_dict(torch.load(model.pt, map_locationcpu)) model.eval() # 切换到推理模式关闭Dropout和BatchNorm的更新 # 读取并预处理图片 image Image.open(test.jpg).convert(RGB) input_tensor transform(image).unsqueeze(0) # 增加batch维度 # 推理 with torch.no_grad(): # 关闭梯度计算节省内存 output model(input_tensor) _, predicted torch.max(output, 1) classes [可回收, 厨余, 有害, 其他] print(f预测结果{classes[predicted.item()]})model.eval()和torch.no_grad()是推理时必加的两行。前者关闭 Dropout 和 BatchNorm 的训练行为后者关闭梯度计算。unsqueeze(0)是因为模型期望的输入是[batch, channel, height, width]单张图片只有后三个维度需要手动补一个 batch 维度。map_locationcpu保证在没有 GPU 的机器上也能加载权重。注意预处理中的Normalize参数必须和训练时一致。如果训练时没用标准化推理时也不要加否则输入分布不匹配预测结果会完全错误。4. 避坑与排查权重加载、类别映射和预处理不一致的常见翻车这个包虽然结构简单但实际跑起来翻车的概率不低。下面是我在帮别人看毕设代码时最常遇到的几个问题按“现象 → 原因 → 解决”整理。4.1 权重加载报错Missing key(s) in state_dict现象运行predict.py时抛出RuntimeError: Error(s) in loading state_dict for SimpleCNN: Missing key(s)...。原因model.py里的网络结构和model.pt保存时的结构不一致。常见情况是作者训练时用了 ResNet但推理时实例化的是 CNN或者改了全连接层的输出维度。解决先打印权重文件的键名和当前模型对比state_dict torch.load(model.pt, map_locationcpu) for k, v in state_dict.items(): print(k, v.shape)对比model.py里定义的层名找到不匹配的地方。如果是全连接层维度不对修改num_classes或fc层的输入维度。如果是层名对不上说明模型结构定义有差异需要根据权重文件的键名反推网络结构。4.2 预测结果全是同一类现象不管输入什么图片输出的类别都一样。原因预处理不一致是最常见的原因。训练时用了Normalize推理时没加或者训练时图片是 BGR 格式推理时用了 RGB。另一个可能是权重文件本身就没训练好。解决先检查predict.py里的transform是否和训练脚本一致。如果找不到训练脚本可以尝试去掉Normalize再跑一次看结果是否变化。如果还是全同一类用torch.load检查权重里的数值是否正常有没有大量零值或 NaN。4.3 CUDA out of memory现象在有 GPU 的机器上跑推理时报显存不足。原因predict.py默认把模型和输入都放到了 GPU 上但显存被其他进程占用了或者图片尺寸太大。解决在加载权重时加map_locationcpu强制用 CPU 推理。或者在推理前加torch.cuda.empty_cache()清理缓存。如果必须用 GPU把输入图片的 batch size 设为 1不要一次加载多张。4.4 类别标签对不上现象模型输出0但classes列表里0对应的类别和实际不符。原因训练时类别的索引顺序和推理时classes列表的顺序不一致。比如训练时0是可回收推理时写成了0是厨余。解决找到训练时的类别映射关系。通常在训练脚本里会有一个class_to_idx的字典或者在数据加载时用ImageFolder会自动按文件夹名排序生成索引。以那个顺序为准不要凭记忆写。4.5 model.pt 和 model1.pt 搞混现象加载model.pt效果正常加载model1.pt效果很差。原因两个权重文件对应不同的模型结构或不同的训练轮数。model1.pt可能是早期保存的检查点或者对应的是 CNN 而不是 ResNet。解决分别用两个权重跑同一张测试图对比输出。如果model1.pt效果差检查它对应的模型结构是否和当前实例化的模型一致。在README.md里通常会说明哪个是最终版本。5. 从能跑到能用替换数据集、调整类别数和验证模型效果的实操技巧把预训练权重跑通只是验证了环境真正要做自己的毕设你得把模型适配到自己的数据集上。这一章讲怎么替换数据、调整类别数以及怎么验证模型是不是真的学到了东西。5.1 用自己的数据重新训练假设你的数据集按文件夹组织每个类别一个文件夹dataset/ ├── recyclable/ ├── kitchen_waste/ ├── hazardous/ └── other/用ImageFolder加载数据然后微调 ResNetfrom torchvision import datasets, transforms from torch.utils.data import DataLoader # 数据增强训练时随机裁剪、翻转增加泛化能力 train_transform transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(dataset/, transformtrain_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue) # 类别数自动从数据集获取 num_classes len(train_dataset.classes) print(f类别{train_dataset.classes}) # 打印索引映射关系ImageFolder会按文件夹名的字母顺序生成索引这个顺序就是推理时classes列表应该用的顺序。batch_size32是常见起点显存不够就降到 16 或 8。shuffleTrue打乱训练顺序防止模型学到数据排列的规律。5.2 训练循环的关键参数import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model get_resnet_model(num_classes).to(device) criterion nn.CrossEntropyLoss() # 只优化fc层的参数学习率设大一点 optimizer optim.Adam(model.fc.parameters(), lr0.001) for epoch in range(10): 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() print(fEpoch {epoch1}, Loss: {running_loss/len(train_loader):.4f}) # 保存权重 torch.save(model.state_dict(), my_model.pt)CrossEntropyLoss是分类任务的标准损失函数内部包含了 softmax。Adam优化器对学习率不敏感lr0.001是安全起点。如果只训练fc层10 个 epoch 通常就能收敛。如果做全网络微调学习率要降到0.0001左右否则预训练权重会被破坏。5.3 验证模型效果的三个检查点训练完不能只看 loss 下降就完事得做实际验证。第一看混淆矩阵。用验证集跑一遍统计每个类别的正确率和误判方向。如果“厨余”大量被误判为“其他”说明这两个类别的特征区分度不够需要增加数据或调整网络。第二看单张图片的置信度。torch.max返回的不仅是类别索引还有置信度分数。如果所有图片的置信度都在 0.3 左右说明模型没学到东西大概率是预处理或标签映射出了问题。第三用训练时没见过的图片测试。从网上找几张垃圾分类的图片手动跑predict.py看输出是否合理。这一步能暴露过拟合问题——训练集准确率很高但新图片全错。5.4 类别数调整的注意事项如果你把四分类改成二分类比如只分可回收和其他需要改两个地方model.py里的num_classes参数以及推理时的classes列表。如果用的是 ResNetmodel.fc nn.Linear(512, num_classes)会自动适配。但要注意如果加载的是原来四分类的权重全连接层的形状不匹配需要重新训练或者只加载卷积层的权重。# 只加载卷积层权重跳过fc层 pretrained_dict torch.load(model.pt) model_dict model.state_dict() # 过滤掉形状不匹配的键 pretrained_dict {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape model_dict[k].shape} model_dict.update(pretrained_dict) model.load_state_dict(model_dict)这段代码的作用是把预训练权重里和当前模型形状匹配的层加载进来不匹配的层比如fc跳过用随机初始化。这样可以在小数据集上快速启动训练。从那以后我每次拿到一个深度学习源码包都强制先跑一遍predict.py确认推理链路通畅再动model.py改结构。顺序反了的话出了问题你分不清是权重的问题还是代码改错了。希望帮到你。本文还有配套的精品资源点击获取
