基于ResNet的垃圾分类识别：.pt模型加载与迁移学习实战

简介基于卷积神经网络 ResNet 算法的垃圾分类识别源码项目内含训练好的预训练模型无需自行训练即可完成图像分类推理。面向计算机科学、人工智能、数据科学等专业学生尤其适合作为课程设计、期末大作业或毕业设计的基础项目帮助快速掌握深度学习图像识别的完整流程。整个资源包一共包含 18 个文件包括 12 张示例图片、两个模型定义与预测源码脚本、两个预训练模型文件以及说明文档等压缩后仅 7.43MB结构清晰、便于直接上手。目前已有 809 人学习代码经过功能验证运行稳定既可直接体验完整推理效果也支持在此基础上进行二次开发丰富分类类别或调整网络结构。压缩包内同时附带验证数据集与说明文档可帮助复现模型预测流程适合用于教学演示或个人学习配套说明文档与示例图片有助于理解数据组织方式、模型加载和预测逻辑对入门进阶和项目立项演示均有实用价值。1. 垃圾分类识别为什么我劝你别从零训练模型手上有基于CNN卷积神经网络resnet算法实现垃圾分类识别源码(含.pt模型).zip想让垃圾分类识别模型快速跑起来、能换自己的数据微调甚至部署到嵌入式或Web端。这类源码包的核心价值不是那几百行训练代码而是里面那个已经收敛好的**.pt模型文件**——它是PyTorch的权重快照意味着你可以直接加载权重做推理不需要显卡也能跑只需要在CPU上装好PyTorch环境。最关键的判断是这个方向该不该投入取决于你手头有没有ResNet预训练模型做迁移学习还是打算从零开始训练一个深残差网络。我从实际工程角度把这条技术路线拆开讲。适合的人是课程设计需要交一个能演示的垃圾分类识别项目、入职后要做目标检测或图像分类的可行性验证、或者想搞懂CNN原理但在真实数据集上跑不动的人。这篇笔记覆盖四件事搞清楚ResNet为什么比普通CNN更适合做分类、如何把zip里的源码组织成一套可复现的训练工程、.pt模型怎么加载和做推理、以及我实测中踩过的那些坑。2. 为什么是ResNet而不是VGG或MobileNet分类任务里的骨架选型2.1 深度加深后ResNet靠“捷径连接”解决退化问题CNN卷积神经网络解决图像分类的基本套路是卷积层提特征、池化层降分辨率、全连接层做决策。但网络一深问题就来了——反而不如浅层网络准这不是过拟合而是梯度在反向传播中逐层衰减前面的层根本学不到东西。ResNet的核心创新是残差块结构上就是给每一层加了一条捷径连接让输出变成F(x) x这样梯度至少能通过捷径直接流回浅层。从代码层面看PyTorch内置的torchvision.models.resnet50已经把这些捷径连接封装好了你不必手写残差块。但要知道一件事垃圾分类识别数据集的单张图片通常包含垃圾袋、饮料瓶这类粗粒度特征而塑料瓶和玻璃瓶的区别往往体现在瓶口螺纹这类细粒度特征上。ResNet的层数越深高层的感受野越大越能融合粗粒度特征和细粒度特征用于分类。50层的resnet50在这个任务上是性价比最高的起点34层偏浅、101层在CPU上推理太慢。2.2 三种常见改造方案从哪一层开始替换全连接层拿到ResNet做迁移学习常见做法是替换最后一层全连接层因为ImageNet的1000类输出和你的垃圾分类类别数对不上。我一般会这样操作import torchvision.models as models import torch.nn as nn def build_model(num_classes, pretrainedTrue): # 用官方预训练权重初始化而不是随机初始化 model models.resnet50(pretrainedpretrained) # 取ResNet最后一层全连接层的输入维度通常是2048 in_features model.fc.in_features # 替换成只输出num_classes个类别的全连接层 model.fc nn.Linear(in_features, num_classes) return model这段代码的逻辑是冻结前面的卷积层参数、只训练新加的全连接层适合小数据集如果你想让模型更贴合垃圾分类场景可以解冻最后几个残差块layer4参与训练学习率调低到主学习率的0.1倍。这里的num_classes要看你的数据集——按常见垃圾种类数设为40或10都行但改了这个数模型的输出维度变了.pt模型的权重文件也要对应重新训练或微调不能拿着别人40类的权重直接改输出维度。提示替换全连接层之后load_state_dict会报“unexpected key”错误解决办法是加载权重时指定strictFalse或者先加载、后替换全连接层。2.3 为什么必须用ImageNet预训练权重而不是随机初始化垃圾分类识别数据集通常只有几千到几万张图片而ResNet50有超过2500万参数。直接随机初始化训练极大概率过拟合验证集准确率卡在70%上下不去。ResNet预训练模型在ImageNet上学到的边缘、纹理、颜色分布等底层特征对垃圾图片同样有效——塑料的纹理、瓶身的反光、纸张的褶皱这些底层视觉模式和ImageNet里的物体表面高度相似。实践上有个误区有人喜欢冻结全部卷积层只训练全连接层但垃圾类别之间的差异很微妙比如“干净纸张”和“纸板箱”的差距可能在颜色而不在纹理。我的做法是分两阶段第一阶段冻结全部卷积层只训练model.fc用较大学习率1e-3快速收敛第二阶段解冻layer3和layer4学习率降到1e-5 ~ 1e-4做全模型微调。这样既能利用预训练特征又能让高层特征适配垃圾分类识别数据的分布。3. 从zip到能训练的工程数据准备、代码结构与训练流程3.1 目录组织和依赖环境拿到zip后第一步做什么这类源码zip解压后常见做法是包含train.py、dataset.py或data_loader.py、model.py、.pt模型文件和requirements.txt。但也有不少zip里只有训练脚本和权重没有完整数据集。所以第一件事不是跑训练而是把目录结构和数据路径理清楚。标准工程结构我建议整理成下面这样garbage_classification/ ├── train.py # 训练入口 ├── inference.py # 推理脚本加载.pt模型做单张图片预测 ├── model.py # 模型定义通常是对resnet的封装 ├── requirements.txt # 依赖清单 ├── data/ │ ├── train/ # 按类别分文件夹存放训练图片 │ └── val/ # 按类别分文件夹存放验证图片 └── weights/ └── best_model.pt # 训练好的权重文件数据目录用ImageFolder格式组织每个子文件夹的名字就是类别名这是PyTorch最省事的读图方式。如果zip里的源码用的是自定义Dataset类路径和标签映射逻辑通常写在dataset.py里需要你根据实际数据路径改一改。3.2 数据集准备用脚本把原始图片整理成ImageFolder格式很多垃圾图片一开始是散落的或者文件名是乱码直接训练会报“Found no valid file for the classes”。我一般会写个小脚本来整理这里用伪代码核心逻辑import os, shutil source_root raw_images # 原始图片目录 target_root data/train # 目标目录 label_map {0: battery, 1: plastic_bottle, 2: paper} # 按实际标签映射 for filename in os.listdir(source_root): if not filename.endswith((.jpg, .jpeg, .png)): continue # 假设文件名前缀是类别ID例如 0_123.jpg label_id filename.split(_)[0] class_dir os.path.join(target_root, label_map[label_id]) os.makedirs(class_dir, exist_okTrue) shutil.copy(os.path.join(source_root, filename), os.path.join(class_dir, filename))逻辑说明label_map是类别ID到类别名的映射这个映射关系必须和训练时的class_to_idx完全一致否则验证时标签对不上。整理完数据后要检查每个类别的图片数量如果某个类只有几张图训练时这个类基本学不出来。注意如果zip里没有数据集只有代码和.pt模型就别纠结造数据了直接用作者训练好的.pt模型做推理更靠谱。训练代码留着微调用。3.3 训练脚本参数配置一组能收敛的基线超参数这是整个工程里最关键的参数配置我的习惯是写一个配置字典或直接用argparse。这里给一组在垃圾分类识别任务上稳定收敛的基线参数import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms from model import build_model # 参数配置 BATCH_SIZE 32 EPOCHS 30 LR 1e-3 NUM_CLASSES 40 DEVICE torch.device(cuda if torch.cuda.is_available() else cpu) # 数据预处理训练集加随机翻转、随机旋转验证集只做归一化 train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(data/train, transformtrain_transform) train_loader DataLoader(train_dataset, batch_sizeBATCH_SIZE, shuffleTrue, num_workers4) model build_model(num_classesNUM_CLASSES, pretrainedTrue).to(DEVICE) optimizer torch.optim.Adam(model.parameters(), lrLR) criterion torch.nn.CrossEntropyLoss() for epoch in range(EPOCHS): 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() # 每个epoch结束后做一次验证代码略 # 保存loss下降时的最优权重而不是最后一个epoch的权重参数说明Resize((224, 224))是ResNet的标准输入尺寸Normalize的参数用ImageNet的均值和标准差因为用了ImageNet预训练权重必须保持一致batch_size在CPU上可调到8或16太大内存顶不住num_workers在Windows上建议设为0否则容易报DataLoader worker错误。损失函数用CrossEntropyLoss语义分割和检测任务里的BCEWithLogitsLoss不能乱换。学习率调度器是后期微调的关键。我会在第15个epoch后把学习率降到1e-4第25个epoch后降到1e-5用torch.optim.lr_scheduler.StepLR或ReduceLROnPlateau都行。如果你在CPU上训练EPOCHS设置5-10就够看趋势了值不值得跑满30轮看前5轮的loss有没有持续下降。3.4 保存模型为什么推荐保存state_dict而不是整个模型zip里给的.pt模型是PyTorch的标准权重格式但PyTorch官方对.pt、.pth、.pth.tar没有语义区分都是torch.save产生的二进制文件。区别在于保存的内容torch.save(model)保存整个模型结构加参数torch.save(model.state_dict())只保存参数。两者都能被load但只保存state_dict是更稳的做法因为模型结构代码一旦变化整包保存的文件就可能彻底废掉。我推荐的保存方式是每个epoch结束做验证验证准确率创新高时保存torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_acc: best_acc, num_classes: NUM_CLASSES }, weights/best_model.pt)逻辑说明把epoch、optimizer_state_dict、best_acc都打包进一个字典这样后续想从断点继续训练断点续训有完整的恢复信息不只是模型权重。num_classes这个键很重要——下次加载时先检查它和当前模型输出维度是否一致。想继续微调而不重新训练时加载这个checkpoint后手动调用optimizer.load_state_dict()就能恢复优化器状态学习率设置会延续之前的调度器状态。4. 加载.pt模型做推理一份能直接抄的CPU代码4.1 推理脚本的完整结构zip里带了.pt模型最终目的通常是做推理演示。这里的核心是推理时的预处理必须和训练时的验证集预处理完全一致。很多人加载模型后准确率暴跌不是模型坏了而是预处理不一致。下面这份代码可以在CPU上直接跑import torch import torchvision.transforms as transforms from PIL import Image from model import build_model # 推理设备选择CPU没有GPU也能跑 device torch.device(cpu) # 类别名称列表顺序必须与训练时class_to_idx完全一致 class_names [battery, plastic_bottle, paper] # 按你的实际类别替换 # 1. 构建模型并加载权重 model build_model(num_classeslen(class_names), pretrainedFalse) checkpoint torch.load(weights/best_model.pt, map_locationdevice) model.load_state_dict(checkpoint[model_state_dict]) model.eval() # 2. 推理预处理Resize和Normalize必须与训练一致 infer_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(image_path): img Image.open(image_path).convert(RGB) tensor infer_transform(img).unsqueeze(0) # 加batch维度 with torch.no_grad(): # 不计算梯度省内存 logits model(tensor) pred_idx torch.argmax(logits, dim1).item() prob torch.softmax(logits, dim1)[0, pred_idx].item() return class_names[pred_idx], prob # 单张图片推理 name, confidence predict(test.jpg) print(f预测类别: {name}, 置信度: {confidence:.2%})逻辑说明pretrainedFalse很关键——推理时不需要加载ImageNet预训练权重而是直接加载你训练好的.pt模型如果这里写True会先下载ImageNet权重再覆盖白白浪费时间和带宽。torch.no_grad()在推理时必须加它让模型不构建计算图显著降低内存占用。.unsqueeze(0)把[3, 224, 224]变成[1, 3, 224, 224]因为ResNet期望输入是四维的批次张量。4.2 .pt模型加载时的四个容易翻车的细节第一个细节map_location必须写。训练时如果用GPU保存的checkpoint权重张量会绑定在cuda:0设备上直接在CPU机器上加载会报no CUDA GPUs are available。torch.load(..., map_locationcpu)强行把权重映射到CPU这是跨设备推理的后悔药。第二个细节model.eval()不能省。不切到eval模式的话BatchNorm层会继续用当前batch的均值和方差在推理时单张图片的统计量极不稳定输出结果会抖动。这也是“加载预训练权重后准确率反而低”的最常见原因——训练代码里model.train()改了状态推理脚本忘了恢复eval()。第三个细节模型定义必须和训练时一致。build_model函数里的num_classes、是否修改了全连接层、是否有Dropout这些结构上的差异会导致load_state_dict报size mismatch。zip里的model.py或models.py就是干这个的直接用它的类定义来build模型别自己另写一份。第四个细节类别列表的顺序。训练时ImageFolder按文件夹名称的字母序生成class_to_idx映射比如battery是0、glass是1推理脚本里的class_names列表必须按这个顺序排不能按中文含义或自己看着顺眼的顺序排。否则预测结果张冠李戴模型实际没坏只是人的映射错了。5. 垃圾分类识别模型避坑指南五条必须知道的实战记录5.1 训练准确率很高、测试集一塌糊涂你的验证集分布和训练集重叠了现象训练30个epoch后训练集准确率98%验证集准确率只有75%而且验证集loss在后期不降反升。原因最常见的是数据泄露——垃圾图片如果按文件名排序分成train和val同一件垃圾的多个角度照片也被拆到了两个集合里。开源的垃圾分类图片经常是同一个饮料瓶从三个角度拍的模型记住了瓶子本身而不是“瓶子”这个类别的通用特征。解决按“物体实例”划分数据而不是按单张图片。没有实例ID标注的话保证一个文件夹里的图片在时间或拍摄序列上不相邻或者做数据增强随机裁剪、翻转、颜色抖动减少模型记住原始像素模式的概率。5.2 加载作者给的.pt模型预测全输出同一个类别现象zip里的best_model.pt加载后对任意图片都预测为“餐巾纸”置信度接近99%。原因标签映射不一致。作者训练时的第5类可能对应“餐巾纸”但你的class_names列表里第5个写的是“易拉罐”于是一切都被映射到错误类别。还有一种可能是作者用torch.save(model)存的整个模型而你用build_model重新build了一个结构加载后虽然没报错但某些层名字对不上权重散落在随机初始化状态。解决先用一张作者训练集里的图片做验证看预测结果是否合理。如果依然全是一个类用torch.load直接看checkpoint里的num_classes和类别数量如果checkpoint是字典格式把里头的model_state_dict的key打印出来和模型对比重点看fc.weight的shape是[40, 2048]还是[10, 2048]。5.3 训练到第10轮loss变成nan模型彻底废了现象loss正常下降几个epoch后突然变成nan之后再也回不来。原因学习率过大导致梯度爆炸。ResNet50用1e-3起步不算激进但解冻layer4做全模型微调时高层的学习率如果也保持1e-3累积梯度很容易把权重推到数值溢出的区间。使用torch.optim.lr_scheduler周期性调整学习率时如果base_lr设置成0.1第一个epoch就可能炸。解决先把torch.clamp或梯度裁剪加上——torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0)。这条在微调阶段比set学习率更管用。另外检查输入图片有没有损坏用PIL.Image.open打开失败的图会得到一个空张量喂进网络后产生nan。数据加载阶段加一个try-except跳过坏图片比训练到一半排查数据省力得多。5.4 CPU上训练慢到怀疑人生batch_size和num_workers的坑现象一台16G内存的笔记本跑ResNet50batch_size设32num_workers设4结果发现训练一步要十几秒内存直接爆了。原因ResNet50的224×224输入尺寸配合batch_size32单batch的中间特征图峰值内存就要8GB以上。num_workers4意味着有4个子进程各自预加载一个batch的图片并做变换内存翻倍。解决CPU训练这个量级的模型我推荐的配置是batch_size8、num_workers0Windows上尤其要设0、图片尺寸降到Resize((160, 160))。ImageFolder的图片Resize在__getitem__阶段执行batch越小CPU压力越小。或者直接放弃训练用curl下载一个官方的分类模型只做推理实训代码只用来做微调——这是最现实的路线。5.5 模型对未见过的垃圾图片识别能力差训练集类别不平衡现象自己拍几张新照片丢进去测试能识别对塑料瓶和易拉罐但纸盒和纸箱经常混。原因数据集里“纸盒”只有几十张而“塑料瓶”有几百张。类别不平衡会让模型倾向于预测样本多的类因为CrossEntropyLoss对多数类的梯度更大。解决给CrossEntropyLoss加weight参数按类别样本数量的倒数归一化让少数类在loss中的占比拉大。代码是weights torch.tensor([1.0 / class_counts[i] for i in range(NUM_CLASSES)]).to(DEVICE)然后nn.CrossEntropyLoss(weightweights)。或者对少数类做更多数据增强——同一张纸盒照片多旋转几个角度、加些椒盐噪声远远比从网上再爬几百张图省事。6. 进阶用法把.pt模型导出ONNX做部署以及验证模型是否真的可信前面把训练和推理跑通了接下来是真正落地时绕不开的一步——部署环境通常没有PyTorch。ONNX是跨平台通用的模型交换格式可以在ONNX Runtime、OpenVINO、TensorRT等推理引擎里运行。转换时要注意PyTorch的BatchNorm层会融合进前面的卷积层所以转换后的ONNX模型大小通常会略小于原始.pt文件。import torch from model import build_model model build_model(num_classeslen(class_names), pretrainedFalse) checkpoint torch.load(weights/best_model.pt, map_locationcpu) model.load_state_dict(checkpoint[model_state_dict]) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, model.onnx, opset_version11, input_names[input], output_names[output] ) print(导出ONNX完成)导出成功后用ONNX Runtime验证一下和PyTorch的推理结果是否一致import onnxruntime as ort import numpy as np session ort.InferenceSession(model.onnx) result session.run([output], {input: dummy_input.numpy()})[0] # 与PyTorch的model(tensor).numpy()结果对比误差应在1e-4以内验证时重点关注两类错误一是预处理差异——ONNX Runtime不吃transforms.Normalize的Tensor操作你需要在C或Python端手动做均值方差归一化二是动态维度——导出时dummy_input是固定[1, 3, 224, 224]如果部署端想支持动态batch_size或动态分辨率需要在torch.onnx.export里加dynamic_axes{input: {0: batch_size}}。丢失PyTorch代码之后ONNX模型才真正具备跨平台复用的能力——这也是那个zip里的.pt模型最终能否进入产线的分水岭。最后提一个我的习惯你验证一个模型时只测测试集准确率是不够的。拿10张家里真实的垃圾照片包括皱了的纸团、半满的饮料瓶、反光的塑料袋单独跑一遍推理。训练集里的图都是规整的、光线均匀的真实场景里但凡室内灯光偏黄或者背景复杂模型的“虚假鲁棒性”就会原形毕露。这个动作比调任何超参数都更能判断这套方案值不值得继续投入。我也是吃过亏之后才养成了这个习惯。希望帮到你。本文还有配套的精品资源点击获取