简介这份资源是一套面向深度学习入门者与环保AI应用开发者的垃圾分类识别项目源码核心解决如何用卷积神经网络自动区分不同类别垃圾的问题适合具备Python基础、希望了解模型训练与跨平台部署流程的学习者。压缩包共6个文件以2个Python脚本和3个CSV数据文件为主另含1个编译缓存文件整体约12KB体量轻便便于快速阅读与二次开发。其中脚本承担模型构建、训练与推理逻辑CSV文件则用于存放类别标签、用户信息与历史记录等配置数据。项目采用ONNX格式导入模型体现了跨框架互操作与灵活部署的思路读者可借此掌握数据预处理、网络结构定义、损失函数与优化器设定、训练循环及模型评估的完整链路并理解模型保存加载与实时图像分类的衔接方式。目前已有116人学习适合作为课程设计或环保类AI项目的参考起点。1. 垃圾分类模型用 ONNX 导入为什么训练完只是走了一半训练脚本跑完best.pt躺在runs/train/exp/weights/里mAP0.5看着挺漂亮很多人到这一步就觉得项目结束了。真到要交付、要嵌进桌面端、要丢给一个只有 CPU 的工控机时才发现 PyTorch 那套环境根本装不进去——CUDA 版本对不上、torch 包 2GB 起步、目标机器连 conda 都没有。这时候 ONNX 就成了那个「后悔药」把模型从训练框架里剥出来变成一个跟框架无关的.onnx文件用 onnxruntime 就能推理。这个标题讲的就是这条链路一个基于深度学习的垃圾分类系统模型不走 PyTorch 原生部署而是导出成 ONNX 再导入运行。它解决的是「训练环境」和「部署环境」解耦的问题适合做毕设、做边缘盒子、做桌面端小工具的人。下面按「先立住原理、再动手复现、最后讲坑」的顺序拆开讲代码都能直接抄。2. 从 PyTorch 到 ONNX垃圾分类模型的导出与结构确认2.1 为什么垃圾分类适合走 ONNX 这条路垃圾分类本质是一个图像分类或检测任务输入是固定尺寸的图片输出是类别概率或检测框。这类模型结构规整、算子常规几乎没有自定义算子是 ONNX 导出最友好的场景之一。常见的做法是用 ResNet、MobileNet、EfficientNet 这类骨干网络做迁移学习或者用 YOLO 系列做检测式分类。选 ONNX 而不是直接上 TensorRT 或 NCNN理由很实际ONNX 是中间格式一次导出可以喂给多个后端。今天用 onnxruntime 在 CPU 上跑明天想换 GPU、换 NCNN、换 RKNN都能从同一个.onnx转过去。热搜里常出现的「onnx转ncnn在线网站」「onnx转rknn」前提都是你手里先有一个干净的 ONNX。所以导出这一步是整个部署链的地基地基歪了后面全歪。需要先明确一件事onnxruntime 和 onnx 不是一回事。onnx 是格式标准定义文件长什么样onnxruntime 是执行引擎负责把文件跑起来。导出用torch.onnx.export推理用onnxruntime.InferenceSession两者别混。2.2 导出前的三个准备动作动手之前先确认三件事否则导出报错会很难查。第一确认模型处于 eval 模式。训练时的 BatchNorm 和 Dropout 在推理时行为不同忘了.eval()会导致导出后的输出和训练时对不上这种玄学问题最耗时间。第二准备一个 dummy input。ONNX 导出需要走一遍前向dummy input 的 shape 必须和真实推理时一致。垃圾分类常见输入是1x3x224x224分类或1x3x640x640检测。第三确认类别数和标签映射。导出只搬模型权重不搬classes.txt。标签顺序要单独存下来推理时按索引查顺序错了模型再准也白搭。import torch import torchvision.models as models # 1. 加载训练好的权重num_classes 换成你自己的类别数 num_classes 6 # 例如可回收物、厨余、有害、其他等 model models.resnet18(weightsNone) model.fc torch.nn.Linear(model.fc.in_features, num_classes) state torch.load(best.pt, map_locationcpu) # 兼容直接存 state_dict 和存 checkpoint 两种情况 state state.get(model, state) if isinstance(state, dict) else state model.load_state_dict(state, strictFalse) # 2. 关键切到推理模式冻结 BN/Dropout model.eval() # 3. dummy input 的 shape 必须和部署时一致 dummy torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy, garbage_cls.onnx, input_names[input], output_names[logits], opset_version12, # 12 兼容性好别盲目追高 dynamic_axes{input: {0: batch}, logits: {0: batch}}, do_constant_foldingTrue, # 常量折叠减小体积 ) print(export done)这段代码的逻辑先重建网络结构结构必须和训练时一致否则load_state_dict会报 key 不匹配再加载权重切 eval然后用一个假输入触发导出。参数说明几个关键点opset_version12是兼容性比较稳的选择opset 太高有些推理后端不认dynamic_axes把 batch 维设成动态这样推理时 batch 可以是 1 也可以是 8不写死do_constant_folding会把能提前算的常量算掉模型更小。strictFalse是给自己留的后路遇到分类头名字对不上时不至于直接崩但导出后一定要验证精度别让它悄悄漏权重。2.3 导出后先验证别急着部署导出成功不等于导出正确。最常见的翻车是ONNX 能跑但输出和 PyTorch 对不上分类结果全乱。所以导出后第一件事是做数值对齐。import numpy as np import onnxruntime as ort # PyTorch 侧输出 with torch.no_grad(): torch_out model(dummy).numpy() # ONNX 侧输出 sess ort.InferenceSession(garbage_cls.onnx, providers[CPUExecutionProvider]) onnx_out sess.run(None, {input: dummy.numpy()})[0] # 比较最大绝对误差一般应小于 1e-4 diff np.max(np.abs(torch_out - onnx_out)) print(max diff:, diff) assert diff 1e-3, 导出误差过大检查 eval() 和 opset逻辑很直白同一个输入分别喂给 PyTorch 和 onnxruntime比输出差。参数上providers指定执行后端CPU 上用CPUExecutionProvider有 GPU 且装了 onnxruntime-gpu 可以换成CUDAExecutionProvider。误差阈值我一般卡1e-4超过1e-3基本说明有问题优先查是不是忘了eval()其次查 opset 是否和某些算子不兼容。这一步过了才说明这个.onnx是可信的。3. 用 onnxruntime 把垃圾分类跑起来预处理、推理、后处理3.1 预处理必须和训练时一模一样模型精度掉得莫名其妙十有八九是预处理不一致。训练时用了什么 resize、什么归一化均值方差推理时必须一字不差。垃圾分类训练常用 ImageNet 的均值和方差mean[0.485,0.456,0.406]std[0.229,0.224,0.225]。如果你训练时用的是自定义的就换成你自己的。from PIL import Image import numpy as np def preprocess(img_path, size224): img Image.open(img_path).convert(RGB) img img.resize((size, size)) # 和训练时的 resize 方式保持一致 x np.asarray(img, dtypenp.float32) / 255.0 mean np.array([0.485, 0.456, 0.406], dtypenp.float32) std np.array([0.229, 0.224, 0.225], dtypenp.float32) x (x - mean) / std x x.transpose(2, 0, 1) # HWC - CHW x np.expand_dims(x, axis0) # 加 batch 维 return np.ascontiguousarray(x, dtypenp.float32)逻辑说明PIL 读进来是 HWC 排列PyTorch/ONNX 要 CHW所以transpose(2,0,1)。np.ascontiguousarray是为了保证内存连续onnxruntime 对非连续数组有时会报奇怪的错。参数上size必须和导出时的 dummy input 对齐导出用 224 推理就得 224除非你导出时把宽高也设成了动态轴。3.2 推理与后处理拿到类别和置信度import onnxruntime as ort import numpy as np CLASSES [可回收物, 厨余垃圾, 有害垃圾, 其他垃圾, 纸张, 塑料] sess ort.InferenceSession(garbage_cls.onnx, providers[CPUExecutionProvider]) input_name sess.get_inputs()[0].name def predict(img_path): x preprocess(img_path) logits sess.run(None, {input_name: x})[0] # shape: [1, num_classes] prob softmax(logits[0]) idx int(np.argmax(prob)) return CLASSES[idx], float(prob[idx]) def softmax(v): v v - np.max(v) # 减最大值防溢出 e np.exp(v) return e / e.sum()逻辑说明sess.get_inputs()[0].name拿到输入节点名导出时我们命名成input但用代码取更稳避免手写错。sess.run返回的是一个列表取[0]是第一个输出。后处理做 softmax 得到概率argmax取最大类别。参数上CLASSES的顺序必须和训练时数据集文件夹的字母序或你定义的映射一致这是最容易错的地方——模型输出索引 3你CLASSES[3]写错结果就全错。3.3 批量推理和耗时测量单张跑通后实际系统往往要处理一批图。ONNX 支持动态 batch前提是导出时设了dynamic_axes。import time def predict_batch(img_paths): xs np.concatenate([preprocess(p) for p in img_paths], axis0) t0 time.time() logits sess.run(None, {input_name: xs})[0] cost (time.time() - t0) / len(img_paths) results [] for row in logits: prob softmax(row) idx int(np.argmax(prob)) results.append((CLASSES[idx], float(prob[idx]))) return results, cost res, per predict_batch([a.jpg, b.jpg, c.jpg]) print(res, f{per*1000:.1f} ms/img)逻辑说明把多张图在 batch 维拼起来一次推理比循环单张快很多因为摊薄了调用开销。per是单张平均耗时用来判断能不能满足实时性要求。参数上batch 不是越大越好CPU 上 batch 太大反而因为内存带宽瓶颈变慢一般 4 到 8 比较合适具体要实测。如果这里报 shape 不匹配说明导出时没设动态 batch回去重导。4. 部署垃圾分类 ONNX 模型时最容易踩的坑4.1 现象ONNX 输出和 PyTorch 差很多分类全乱原因最常见是导出时忘了model.eval()BatchNorm 还在用 batch 统计量其次是 opset 版本和某个算子不兼容导出时被静默替换成了近似实现。解决导出前强制model.eval()导出后立刻做 2.3 节的数值对齐误差超1e-3就换 opset12 不行试 11 或 13。别跳过对齐直接部署那是给自己埋雷。4.2 现象推理时报输入 shape 不匹配原因导出时 dummy input 写死了1x3x224x224且没设dynamic_axes推理时传了别的尺寸或 batch。解决要么推理时严格按导出尺寸来要么重新导出并加上dynamic_axes。注意动态轴只对你声明的那一维生效宽高要动态就得把 2、3 维也写进去。4.3 现象CPU 上单张要几百毫秒根本没法用原因用了 ResNet50 这种大骨干或者没开图优化或者线程数没配。解决换 MobileNetV3、ShuffleNet 这类轻量骨干重训创建 session 时配sess_options把intra_op_num_threads设成物理核数有条件就上onnxruntime-gpu。热搜里的「.onnx量化int8」也是条路量化后体积和耗时都能降但精度会掉一点要重新验证。so ort.SessionOptions() so.intra_op_num_threads 4 # 按物理核数设 so.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess ort.InferenceSession(garbage_cls.onnx, so, providers[CPUExecutionProvider])4.4 现象换台机器结果就变了原因预处理里的均值方差、resize 插值方式、甚至 PIL 版本不同导致的解码差异都会让输入像素值有微小偏移累积到分类边界样本上就翻类。解决把预处理参数写进配置文件和模型一起交付对边界样本置信度 0.5 附近做二次确认或直接判为「不确定」别硬给一个类别。4.5 现象类别索引对不上A 类被判成 B 类原因训练时ImageFolder按文件夹名字母序生成索引部署时CLASSES列表手写顺序不一致。解决训练完立刻把dataset.classes存成labels.json部署时读这个文件永远不要手敲类别顺序。这是血泪经验手敲迟早出错。5. 让垃圾分类 ONNX 更实用量化、验证与一个交付习惯模型能跑之后下一步通常是压体积、提速度。INT8 量化是最直接的手段onnxruntime 自带动态量化几行代码就能把模型砍掉一半以上。from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_inputgarbage_cls.onnx, model_outputgarbage_cls_int8.onnx, weight_typeQuantType.QInt8, ) print(quantized)逻辑说明动态量化只量化权重激活值在推理时动态量化不需要校准数据集最省事。参数上weight_typeQInt8是 8 位整型体积约为 FP32 的四分之一。代价是精度可能掉 1 到 3 个百分点所以量化后必须重跑验证集对比量化前后的准确率。如果掉太多就改用静态量化并准备校准集或者只量化部分层。量化完别只看文件大小要拿一批真实图片跑一遍统计准确率和单张耗时做个前后对比表版本体积单张耗时(CPU)验证集准确率FP32 ONNX基准基准基准INT8 动态量化约 1/4通常降 30%~50%可能降 1~3 个点这张表要自己实测填别抄别人的数字硬件不同差异很大。我一般会留一个benchmark.py每次换模型或换量化策略都跑一遍数据说话。最后说个交付习惯把.onnx、labels.json、预处理参数、一个最小推理示例脚本打包在一起写清楚输入尺寸和均值方差。我吃过亏——模型交付出去对方预处理用了[0,0,0]均值结果准确率惨不忍睹来回扯皮好几天。从那以后凡是交付 ONNX我一定附一个能直接跑的infer.py让对方先跑通再谈集成。垃圾分类这种项目模型只是中间产物能稳定跑起来、结果可复现才算真正做完。希望帮到你。本文还有配套的精品资源点击获取