尧图网页设计 
简介这份资源是一套面向深度学习入门与计算机视觉实践者的垃圾分类识别项目源码适合希望了解图像分类完整流程、并尝试将模型导出为ONNX格式进行跨平台部署的开发者。压缩包共6个文件以Python脚本、CSV数据文件和编译缓存文件为主整体约12KB体量轻巧便于快速阅读与二次开发。其中Python代码承担模型构建、训练与推理逻辑CSV文件用于存放类别标签、用户信息与历史记录等配置数据pyc文件为编译后的缓存。项目围绕卷积神经网络对垃圾图像进行特征提取与类别预测并借助ONNX实现不同框架间的模型互操作提升部署灵活性。代码结构涵盖数据预处理、模型定义、训练循环、评估及模型导入导出等环节目录组织清晰可作为课程设计、毕业设计或自学练手的参考模板。目前已有116人学习下载适合想快速跑通垃圾分类demo、理解ONNX导入流程的读者参考借鉴。1. 垃圾分类模型用 ONNX 导入为什么这条路线值得走训练一个垃圾分类模型不难难的是把它塞进别人的系统里跑起来。你大概率遇到过这种场景自己用 PyTorch 训好了 ResNet 或者 MobileNet准确率看着还行结果对方一句“我们这边是 C 环境”或者“部署机器上装不了 PyTorch”整个交付就卡住了。基于深度学习的垃圾分类系统模型使用 ONNX 导入解决的正是这个“最后一公里”的问题——把训练框架和推理框架解耦让模型变成一个通用的中间格式谁都能加载。这条路线适合两类人一类是做深度学习毕设或者实战项目案例的学生需要把模型交给一个不依赖训练环境的演示程序另一类是做边缘部署的工程师目标平台可能是 ONNX Runtime、TensorRT甚至是转成 RKNN 或 NCNN 跑在 NPU 上。核心逻辑就一句话PyTorch 负责训ONNX 负责当“中间人”推理端只认 ONNX。下面从导出、验证到落地把这条链路拆开讲。2. 从 PyTorch 到 ONNX垃圾分类模型的导出与验证2.1 为什么垃圾分类模型适合走 ONNX 这条路垃圾分类本质上是一个图像分类任务输入是固定尺寸的图片输出是几个类别概率。这类模型结构规整没有复杂的动态控制流导出 ONNX 的成功率非常高。常见的 backbone 比如 ResNet50、MobileNetV3、EfficientNet-B0在 torchvision 里都有预训练权重微调之后直接导出即可。选 ONNX 而不是直接保存state_dict理由有三个。第一ONNX 是开放标准ONNX Runtime 在 Windows、Linux、ARM 上都有预编译包装起来比 PyTorch 轻得多。第二ONNX 可以进一步转成 TensorRT、OpenVINO、NCNN、RKNN覆盖从服务器到边缘设备的全部场景。第三ONNX 文件自带计算图定义推理端不需要知道模型是用什么框架训的也不需要复现网络结构代码。我一般会建议只要你的模型没有自定义算子、没有动态 shape 的硬需求导出 ONNX 就是默认动作。垃圾分类这种任务输入尺寸固定为 224×224 或 320×320batch size 在推理时通常设为 1完全符合 ONNX 的舒适区。2.2 导出 ONNX 的最小可复现脚本下面这段代码假设你已经用 PyTorch 训好了一个垃圾分类模型保存在best_model.pth网络结构定义在model.py里。导出时最关键的是dummy_input的 shape 要和推理时一致opset_version建议用 11 或 12兼容性最好。import torch import torch.onnx from model import GarbageClassifier # 你定义的网络结构 # 1. 加载训练好的权重 model GarbageClassifier(num_classes6) # 假设分6类可回收、厨余、有害、其他等 checkpoint torch.load(best_model.pth, map_locationcpu) model.load_state_dict(checkpoint[model_state_dict]) model.eval) # 必须切到推理模式否则BN和Dropout行为不对 # 2. 构造虚拟输入shape必须和推理时完全一致 dummy_input torch.randn(1, 3, 224, 224) # 3. 导出ONNX torch.onnx.export( model, dummy_input, garbage_classifier.onnx, export_paramsTrue, # 把权重也写进文件 opset_version11, # 算子集版本11兼容性好 do_constant_foldingTrue, # 常量折叠优化 input_names[input], # 推理端按这个名字喂数据 output_names[output], # 推理端按这个名字取结果 dynamic_axesNone # 垃圾分类推理固定batch1不需要动态轴 ) print(ONNX导出完成)这段代码里最容易翻车的地方是model.eval()。训练时 BatchNorm 层用的是当前 batch 的均值和方差推理时必须用训练时累积的滑动平均值。如果忘了加这一行导出的 ONNX 模型在推理时输出会完全乱掉而且这种错误不会报异常属于典型的“玄学”问题。另一个参数是opset_version。版本太低可能不支持某些算子版本太高则老版本的 ONNX Runtime 加载不了。我一般用 11如果模型里有Hardswish这类较新的激活函数可能需要升到 12 或 13。导出时如果报错说某个算子不支持优先查这个参数。dynamic_axesNone表示输入输出都是固定 shape。垃圾分类推理通常一次处理一张图固定 batch1 最省事。如果你确实需要动态 batch可以改成{input: {0: batch_size}, output: {0: batch_size}}但要注意后续推理端也要对应处理。2.3 导出后的验证别跳过这一步导出完成不等于正确。我见过太多次导出成功但推理结果和 PyTorch 对不上的情况。验证方法很简单用同一张图片分别跑 PyTorch 和 ONNX Runtime比较输出差异。import numpy as np import onnxruntime as ort import torch from PIL import Image from torchvision import transforms # 1. 准备一张测试图片 img Image.open(test.jpg).convert(RGB) preprocess transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) input_tensor preprocess(img).unsqueeze(0) # shape: [1,3,224,224] # 2. PyTorch推理 model.eval() with torch.no_grad(): torch_output model(input_tensor).numpy() # 3. ONNX Runtime推理 session ort.InferenceSession(garbage_classifier.onnx) onnx_output session.run( [output], {input: input_tensor.numpy()} )[0] # 4. 比较差异 diff np.abs(torch_output - onnx_output).max() print(f最大绝对误差: {diff:.6f}) print(fPyTorch预测类别: {torch_output.argmax()}) print(fONNX预测类别: {onnx_output.argmax()})判断标准最大绝对误差在 1e-4 以内算正常1e-3 以内可以接受超过 1e-2 说明导出有问题。如果误差大先检查model.eval()是否加了再检查预处理是否完全一致。预处理不一致是第二大坑——比如 PyTorch 用了 NormalizeONNX 推理时忘了做同样的归一化结果自然对不上。提示验证时用的图片最好从训练集和验证集里各取几张确保模型对已知类别和未知类别的输出都一致。3. ONNX Runtime 推理部署从加载到出结果3.1 ONNX Runtime 和 ONNX 的区别先搞清楚很多人会把 ONNX 和 ONNX Runtime 混为一谈。ONNX 是模型格式标准定义的是文件长什么样ONNX Runtime 是微软开源的推理引擎负责把 ONNX 文件加载进来、在 CPU 或 GPU 上执行计算。类比一下ONNX 是 PDF 格式ONNX Runtime 是 PDF 阅读器。你可以用别的阅读器打开 PDF也可以用别的引擎跑 ONNX比如 TensorRT、OpenVINO。在垃圾分类系统里ONNX Runtime 是最省心的选择。pip 直接装CPU 版本不需要 CUDAWindows 和 Linux 都有 wheel 包。如果你的部署机器有 NVIDIA 显卡可以装onnxruntime-gpu推理速度能快 5 到 10 倍。3.2 一个完整的垃圾分类推理函数下面这个函数把图片预处理、ONNX 推理、结果后处理串起来可以直接放进你的 Flask 或者 FastAPI 服务里。import numpy as np import onnxruntime as ort from PIL import Image from torchvision import transforms # 类别名称顺序必须和训练时一致 CLASS_NAMES [可回收物, 厨余垃圾, 有害垃圾, 其他垃圾, 纸张, 塑料] # 初始化session只做一次 session ort.InferenceSession( garbage_classifier.onnx, providers[CPUExecutionProvider] # 有GPU就换成CUDAExecutionProvider ) # 预处理管道必须和训练时完全一致 preprocess transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def classify_garbage(image_path): 输入图片路径返回类别名和置信度 img Image.open(image_path).convert(RGB) input_tensor preprocess(img).unsqueeze(0).numpy() # 推理 outputs session.run( [output], {input: input_tensor} )[0] # shape: [1, num_classes] # softmax转概率 exp np.exp(outputs - np.max(outputs, axis1, keepdimsTrue)) probs exp / np.sum(exp, axis1, keepdimsTrue) pred_idx int(np.argmax(probs, axis1)[0]) confidence float(probs[0][pred_idx]) return CLASS_NAMES[pred_idx], confidence # 调用示例 label, conf classify_garbage(test.jpg) print(f分类结果: {label}, 置信度: {conf:.4f})几个关键点说明。providers参数决定用什么后端跑CPU 版填CPUExecutionProviderGPU 版填CUDAExecutionProvider。如果填了 GPU 但环境没配好ONNX Runtime 会自动回退到 CPU不会报错但速度上不去这点要注意。预处理管道必须和训练时一模一样。Resize(256)然后CenterCrop(224)是 torchvision 的经典组合如果你的训练代码用的是Resize(224)直接缩放这里也要改成Resize(224)。Normalize 的均值和标准差也是固定的 ImageNet 统计值除非你训练时用了别的。后处理里的 softmax 不是必须的因为很多分类模型在导出时已经把 softmax 包含在计算图里了。你可以先跑一次看看输出值是否在 0 到 1 之间且和为 1如果是就不需要再做 softmax。我一般会在导出时把 softmax 留在模型里推理端直接取 argmax 就行。3.3 性能调优让推理快起来的几个参数ONNX Runtime 默认配置已经不错但有几个参数可以调。在创建InferenceSession时传入sess_optionsimport onnxruntime as ort sess_options ort.SessionOptions() sess_options.intra_op_num_threads 4 # 单算子内部并行线程数 sess_options.inter_op_num_threads 2 # 算子之间并行线程数 sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL session ort.InferenceSession( garbage_classifier.onnx, sess_optionssess_options, providers[CPUExecutionProvider] )intra_op_num_threads控制单个算子比如矩阵乘法用几个线程一般设为物理核心数。inter_op_num_threads控制不同算子之间的并行度设太大反而会因为线程切换开销导致变慢2 到 4 比较合适。graph_optimization_level设为ORT_ENABLE_ALL会启用算子融合、常量折叠等优化推理速度通常能提升 10% 到 20%。如果你的部署环境是 ARM 开发板或者手机CPU 核心数少建议把intra_op_num_threads设为 2 或 1避免线程争抢。这些参数没有万能值最好在目标机器上实测几组。注意onnxruntime-gpu和onnxruntime不能同时安装会冲突。装 GPU 版之前先pip uninstall onnxruntime。4. 垃圾分类 ONNX 模型部署的避坑与排查4.1 导出时报 “Unsupported operator”现象torch.onnx.export执行到一半抛出异常提示某个算子不支持。原因PyTorch 的某些算子在高版本 opset 才被支持或者你用了自定义算子。垃圾分类模型里常见的是HardswishMobileNetV3 用和AdaptiveAvgPool2d的某些变体。解决先把opset_version升到 12 或 13 试试。如果还不行检查模型里有没有自己写的forward函数用了 Python 控制流if/forONNX 导出时会尝试追踪计算图遇到数据依赖的控制流就会失败。解决办法是把控制流改成固定结构或者用torch.jit.script先转 TorchScript 再导出。4.2 推理结果和 PyTorch 对不上现象ONNX 推理输出的类别和 PyTorch 完全不一样或者置信度差异很大。原因最常见的是忘了model.eval()其次是预处理不一致第三是输入 tensor 的 layout 不对。PyTorch 默认是 NCHWONNX 也是 NCHW但如果你在预处理时用了transforms.ToTensor()之后又做了 transpose就会变成 NHWC导致推理结果错乱。解决按顺序排查。先确认model.eval()加了再逐行对比预处理代码最后打印输入 tensor 的 shape 确认是[1, 3, 224, 224]。如果都没问题用 2.3 节的验证脚本跑一遍看误差有多大。4.3 ONNX Runtime 加载模型时报 shape 不匹配现象session.run时抛出InvalidArgument: Got invalid dimensions for input。原因推理时喂进去的 tensor shape 和导出时定义的dummy_inputshape 不一致。比如导出时是[1, 3, 224, 224]推理时图片预处理后变成了[1, 3, 256, 256]。解决检查预处理管道里的Resize和CenterCrop尺寸确保最终输出是 224×224。如果确实需要支持多种输入尺寸导出时把dynamic_axes里的高度和宽度也设为动态但这样会失去一些优化机会推理速度会下降。4.4 GPU 推理比 CPU 还慢现象装了onnxruntime-gpu但推理速度没有提升甚至更慢。原因模型太小数据传输开销超过了计算节省的时间。垃圾分类用的 MobileNet 或 ResNet18 参数量少GPU 的 kernel launch 开销占比高。另外如果每次推理都新建 sessionGPU 初始化时间也会拖慢整体速度。解决先确认providers里CUDAExecutionProvider排在第一位。然后把 session 做成全局变量只初始化一次。如果模型确实小CPU 推理反而更划算。我实测过 MobileNetV3 在 CPU 上单张推理 15ms 左右GPU 上 8ms但加上数据传输和初始化批量处理时 GPU 才有明显优势。4.5 转 RKNN 或 NCNN 后精度下降现象ONNX 模型在 ONNX Runtime 上精度正常转成 RKNN 或 NCNN 后精度明显下降。原因量化。RKNN 和 NCNN 默认会做 int8 量化量化过程会引入误差。垃圾分类模型如果对某些类别区分度不高量化后容易混淆。解决先确认转换时是否开启了量化。如果开启了尝试关闭量化或者用混合量化只量化部分层。另外转换工具对 ONNX 的 opset 版本有要求RKNN 通常要求 opset 11 或 12NCNN 对 11 支持最好。转换前先用onnxsim简化模型去掉冗余算子能提高转换成功率。5. 把 ONNX 垃圾分类模型塞进边缘设备的几个技巧5.1 用 onnxsim 给模型“瘦身”导出后的 ONNX 文件里往往有很多冗余算子比如连续的 Reshape、Transpose或者常量折叠没做干净的节点。onnxsim可以自动简化计算图减小文件体积有时还能提升推理速度。pip install onnxsim onnxsim garbage_classifier.onnx garbage_classifier_sim.onnx执行后会输出简化前后的算子数量对比。我遇到过导出后 120 个算子简化后剩 80 个的情况文件从 45MB 降到 42MB。简化后的模型精度不变因为 onnxsim 只做等价变换。5.2 转 RKNN 的完整命令与参数如果你要把模型跑在瑞芯微的 NPU 上需要转成 RKNN 格式。转换工具是rknn-toolkit2在 x86 机器上运行。from rknn.api import RKNN rknn RKNN() rknn.config( mean_values[[0.485 * 255, 0.456 * 255, 0.406 * 255]], std_values[[0.229 * 255, 0.224 * 255, 0.225 * 255]], target_platformrk3588 ) rknn.load_onnx(modelgarbage_classifier_sim.onnx) rknn.build(do_quantizationTrue, datasetquant_dataset.txt) rknn.export_rknn(garbage_classifier.rknn)mean_values和std_values要和训练时的 Normalize 对应注意这里乘了 255因为 RKNN 内部处理的是 0 到 255 的像素值。do_quantizationTrue会做 int8 量化需要提供一个quant_dataset.txt里面是几十张训练图片的路径用于校准量化参数。量化后模型体积缩小到原来的四分之一左右推理速度提升明显但精度可能掉 1 到 3 个百分点。如果精度掉得太多把do_quantization设为False用 float16 推理精度基本无损速度也有提升。5.3 一个容易被忽略的细节类别顺序ONNX 模型输出的类别顺序取决于你训练时CLASS_NAMES的定义顺序。这个顺序在导出、推理、转 RKNN 的整个链路里必须保持一致。我见过有人在训练时用[可回收, 厨余, 有害, 其他]推理时写成了[厨余, 可回收, 有害, 其他]结果所有分类都错位了。建议把类别名称写在一个单独的 JSON 文件里训练和推理都从这个文件读。这样改的时候只改一处不会漏。{ class_names: [可回收物, 厨余垃圾, 有害垃圾, 其他垃圾, 纸张, 塑料] }推理代码里用json.load读进来不要硬编码在 Python 文件里。这个习惯能帮你省下不少排查时间。5.4 验证边缘设备推理结果的方法模型转到 RKNN 或 NCNN 之后怎么确认结果是对的最直接的方法是在 PC 上用 ONNX Runtime 跑一遍记录每张测试图的输出向量然后在边缘设备上跑同样的图比较两个输出向量的余弦相似度。相似度在 0.99 以上算正常低于 0.95 说明转换过程有问题。如果没有条件在边缘设备上逐张对比至少要在转换后用 RKNN 的模拟器跑一遍验证集看准确率掉了多少。掉 1 到 2 个点可以接受掉 5 个点以上就要检查量化配置和预处理参数了。我自己做垃圾分类项目时习惯在导出 ONNX 之后先不急着转其他格式而是用 ONNX Runtime 跑完整验证集确认精度和 PyTorch 一致。这一步过了后面转 RKNN 或者 TensorRT 才有意义。如果 ONNX 这关就没过后面全是白费功夫。希望帮到你。本文还有配套的精品资源点击获取 基于ONNX的垃圾分类模型部署：从PyTorch导出到边缘推理实战
简介这份资源是一套面向深度学习入门与计算机视觉实践者的垃圾分类识别项目源码适合希望了解图像分类完整流程、并尝试将模型导出为ONNX格式进行跨平台部署的开发者。压缩包共6个文件以Python脚本、CSV数据文件和编译缓存文件为主整体约12KB体量轻巧便于快速阅读与二次开发。其中Python代码承担模型构建、训练与推理逻辑CSV文件用于存放类别标签、用户信息与历史记录等配置数据pyc文件为编译后的缓存。项目围绕卷积神经网络对垃圾图像进行特征提取与类别预测并借助ONNX实现不同框架间的模型互操作提升部署灵活性。代码结构涵盖数据预处理、模型定义、训练循环、评估及模型导入导出等环节目录组织清晰可作为课程设计、毕业设计或自学练手的参考模板。目前已有116人学习下载适合想快速跑通垃圾分类demo、理解ONNX导入流程的读者参考借鉴。1. 垃圾分类模型用 ONNX 导入为什么这条路线值得走训练一个垃圾分类模型不难难的是把它塞进别人的系统里跑起来。你大概率遇到过这种场景自己用 PyTorch 训好了 ResNet 或者 MobileNet准确率看着还行结果对方一句“我们这边是 C 环境”或者“部署机器上装不了 PyTorch”整个交付就卡住了。基于深度学习的垃圾分类系统模型使用 ONNX 导入解决的正是这个“最后一公里”的问题——把训练框架和推理框架解耦让模型变成一个通用的中间格式谁都能加载。这条路线适合两类人一类是做深度学习毕设或者实战项目案例的学生需要把模型交给一个不依赖训练环境的演示程序另一类是做边缘部署的工程师目标平台可能是 ONNX Runtime、TensorRT甚至是转成 RKNN 或 NCNN 跑在 NPU 上。核心逻辑就一句话PyTorch 负责训ONNX 负责当“中间人”推理端只认 ONNX。下面从导出、验证到落地把这条链路拆开讲。2. 从 PyTorch 到 ONNX垃圾分类模型的导出与验证2.1 为什么垃圾分类模型适合走 ONNX 这条路垃圾分类本质上是一个图像分类任务输入是固定尺寸的图片输出是几个类别概率。这类模型结构规整没有复杂的动态控制流导出 ONNX 的成功率非常高。常见的 backbone 比如 ResNet50、MobileNetV3、EfficientNet-B0在 torchvision 里都有预训练权重微调之后直接导出即可。选 ONNX 而不是直接保存state_dict理由有三个。第一ONNX 是开放标准ONNX Runtime 在 Windows、Linux、ARM 上都有预编译包装起来比 PyTorch 轻得多。第二ONNX 可以进一步转成 TensorRT、OpenVINO、NCNN、RKNN覆盖从服务器到边缘设备的全部场景。第三ONNX 文件自带计算图定义推理端不需要知道模型是用什么框架训的也不需要复现网络结构代码。我一般会建议只要你的模型没有自定义算子、没有动态 shape 的硬需求导出 ONNX 就是默认动作。垃圾分类这种任务输入尺寸固定为 224×224 或 320×320batch size 在推理时通常设为 1完全符合 ONNX 的舒适区。2.2 导出 ONNX 的最小可复现脚本下面这段代码假设你已经用 PyTorch 训好了一个垃圾分类模型保存在best_model.pth网络结构定义在model.py里。导出时最关键的是dummy_input的 shape 要和推理时一致opset_version建议用 11 或 12兼容性最好。import torch import torch.onnx from model import GarbageClassifier # 你定义的网络结构 # 1. 加载训练好的权重 model GarbageClassifier(num_classes6) # 假设分6类可回收、厨余、有害、其他等 checkpoint torch.load(best_model.pth, map_locationcpu) model.load_state_dict(checkpoint[model_state_dict]) model.eval) # 必须切到推理模式否则BN和Dropout行为不对 # 2. 构造虚拟输入shape必须和推理时完全一致 dummy_input torch.randn(1, 3, 224, 224) # 3. 导出ONNX torch.onnx.export( model, dummy_input, garbage_classifier.onnx, export_paramsTrue, # 把权重也写进文件 opset_version11, # 算子集版本11兼容性好 do_constant_foldingTrue, # 常量折叠优化 input_names[input], # 推理端按这个名字喂数据 output_names[output], # 推理端按这个名字取结果 dynamic_axesNone # 垃圾分类推理固定batch1不需要动态轴 ) print(ONNX导出完成)这段代码里最容易翻车的地方是model.eval()。训练时 BatchNorm 层用的是当前 batch 的均值和方差推理时必须用训练时累积的滑动平均值。如果忘了加这一行导出的 ONNX 模型在推理时输出会完全乱掉而且这种错误不会报异常属于典型的“玄学”问题。另一个参数是opset_version。版本太低可能不支持某些算子版本太高则老版本的 ONNX Runtime 加载不了。我一般用 11如果模型里有Hardswish这类较新的激活函数可能需要升到 12 或 13。导出时如果报错说某个算子不支持优先查这个参数。dynamic_axesNone表示输入输出都是固定 shape。垃圾分类推理通常一次处理一张图固定 batch1 最省事。如果你确实需要动态 batch可以改成{input: {0: batch_size}, output: {0: batch_size}}但要注意后续推理端也要对应处理。2.3 导出后的验证别跳过这一步导出完成不等于正确。我见过太多次导出成功但推理结果和 PyTorch 对不上的情况。验证方法很简单用同一张图片分别跑 PyTorch 和 ONNX Runtime比较输出差异。import numpy as np import onnxruntime as ort import torch from PIL import Image from torchvision import transforms # 1. 准备一张测试图片 img Image.open(test.jpg).convert(RGB) preprocess transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) input_tensor preprocess(img).unsqueeze(0) # shape: [1,3,224,224] # 2. PyTorch推理 model.eval() with torch.no_grad(): torch_output model(input_tensor).numpy() # 3. ONNX Runtime推理 session ort.InferenceSession(garbage_classifier.onnx) onnx_output session.run( [output], {input: input_tensor.numpy()} )[0] # 4. 比较差异 diff np.abs(torch_output - onnx_output).max() print(f最大绝对误差: {diff:.6f}) print(fPyTorch预测类别: {torch_output.argmax()}) print(fONNX预测类别: {onnx_output.argmax()})判断标准最大绝对误差在 1e-4 以内算正常1e-3 以内可以接受超过 1e-2 说明导出有问题。如果误差大先检查model.eval()是否加了再检查预处理是否完全一致。预处理不一致是第二大坑——比如 PyTorch 用了 NormalizeONNX 推理时忘了做同样的归一化结果自然对不上。提示验证时用的图片最好从训练集和验证集里各取几张确保模型对已知类别和未知类别的输出都一致。3. ONNX Runtime 推理部署从加载到出结果3.1 ONNX Runtime 和 ONNX 的区别先搞清楚很多人会把 ONNX 和 ONNX Runtime 混为一谈。ONNX 是模型格式标准定义的是文件长什么样ONNX Runtime 是微软开源的推理引擎负责把 ONNX 文件加载进来、在 CPU 或 GPU 上执行计算。类比一下ONNX 是 PDF 格式ONNX Runtime 是 PDF 阅读器。你可以用别的阅读器打开 PDF也可以用别的引擎跑 ONNX比如 TensorRT、OpenVINO。在垃圾分类系统里ONNX Runtime 是最省心的选择。pip 直接装CPU 版本不需要 CUDAWindows 和 Linux 都有 wheel 包。如果你的部署机器有 NVIDIA 显卡可以装onnxruntime-gpu推理速度能快 5 到 10 倍。3.2 一个完整的垃圾分类推理函数下面这个函数把图片预处理、ONNX 推理、结果后处理串起来可以直接放进你的 Flask 或者 FastAPI 服务里。import numpy as np import onnxruntime as ort from PIL import Image from torchvision import transforms # 类别名称顺序必须和训练时一致 CLASS_NAMES [可回收物, 厨余垃圾, 有害垃圾, 其他垃圾, 纸张, 塑料] # 初始化session只做一次 session ort.InferenceSession( garbage_classifier.onnx, providers[CPUExecutionProvider] # 有GPU就换成CUDAExecutionProvider ) # 预处理管道必须和训练时完全一致 preprocess transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def classify_garbage(image_path): 输入图片路径返回类别名和置信度 img Image.open(image_path).convert(RGB) input_tensor preprocess(img).unsqueeze(0).numpy() # 推理 outputs session.run( [output], {input: input_tensor} )[0] # shape: [1, num_classes] # softmax转概率 exp np.exp(outputs - np.max(outputs, axis1, keepdimsTrue)) probs exp / np.sum(exp, axis1, keepdimsTrue) pred_idx int(np.argmax(probs, axis1)[0]) confidence float(probs[0][pred_idx]) return CLASS_NAMES[pred_idx], confidence # 调用示例 label, conf classify_garbage(test.jpg) print(f分类结果: {label}, 置信度: {conf:.4f})几个关键点说明。providers参数决定用什么后端跑CPU 版填CPUExecutionProviderGPU 版填CUDAExecutionProvider。如果填了 GPU 但环境没配好ONNX Runtime 会自动回退到 CPU不会报错但速度上不去这点要注意。预处理管道必须和训练时一模一样。Resize(256)然后CenterCrop(224)是 torchvision 的经典组合如果你的训练代码用的是Resize(224)直接缩放这里也要改成Resize(224)。Normalize 的均值和标准差也是固定的 ImageNet 统计值除非你训练时用了别的。后处理里的 softmax 不是必须的因为很多分类模型在导出时已经把 softmax 包含在计算图里了。你可以先跑一次看看输出值是否在 0 到 1 之间且和为 1如果是就不需要再做 softmax。我一般会在导出时把 softmax 留在模型里推理端直接取 argmax 就行。3.3 性能调优让推理快起来的几个参数ONNX Runtime 默认配置已经不错但有几个参数可以调。在创建InferenceSession时传入sess_optionsimport onnxruntime as ort sess_options ort.SessionOptions() sess_options.intra_op_num_threads 4 # 单算子内部并行线程数 sess_options.inter_op_num_threads 2 # 算子之间并行线程数 sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL session ort.InferenceSession( garbage_classifier.onnx, sess_optionssess_options, providers[CPUExecutionProvider] )intra_op_num_threads控制单个算子比如矩阵乘法用几个线程一般设为物理核心数。inter_op_num_threads控制不同算子之间的并行度设太大反而会因为线程切换开销导致变慢2 到 4 比较合适。graph_optimization_level设为ORT_ENABLE_ALL会启用算子融合、常量折叠等优化推理速度通常能提升 10% 到 20%。如果你的部署环境是 ARM 开发板或者手机CPU 核心数少建议把intra_op_num_threads设为 2 或 1避免线程争抢。这些参数没有万能值最好在目标机器上实测几组。注意onnxruntime-gpu和onnxruntime不能同时安装会冲突。装 GPU 版之前先pip uninstall onnxruntime。4. 垃圾分类 ONNX 模型部署的避坑与排查4.1 导出时报 “Unsupported operator”现象torch.onnx.export执行到一半抛出异常提示某个算子不支持。原因PyTorch 的某些算子在高版本 opset 才被支持或者你用了自定义算子。垃圾分类模型里常见的是HardswishMobileNetV3 用和AdaptiveAvgPool2d的某些变体。解决先把opset_version升到 12 或 13 试试。如果还不行检查模型里有没有自己写的forward函数用了 Python 控制流if/forONNX 导出时会尝试追踪计算图遇到数据依赖的控制流就会失败。解决办法是把控制流改成固定结构或者用torch.jit.script先转 TorchScript 再导出。4.2 推理结果和 PyTorch 对不上现象ONNX 推理输出的类别和 PyTorch 完全不一样或者置信度差异很大。原因最常见的是忘了model.eval()其次是预处理不一致第三是输入 tensor 的 layout 不对。PyTorch 默认是 NCHWONNX 也是 NCHW但如果你在预处理时用了transforms.ToTensor()之后又做了 transpose就会变成 NHWC导致推理结果错乱。解决按顺序排查。先确认model.eval()加了再逐行对比预处理代码最后打印输入 tensor 的 shape 确认是[1, 3, 224, 224]。如果都没问题用 2.3 节的验证脚本跑一遍看误差有多大。4.3 ONNX Runtime 加载模型时报 shape 不匹配现象session.run时抛出InvalidArgument: Got invalid dimensions for input。原因推理时喂进去的 tensor shape 和导出时定义的dummy_inputshape 不一致。比如导出时是[1, 3, 224, 224]推理时图片预处理后变成了[1, 3, 256, 256]。解决检查预处理管道里的Resize和CenterCrop尺寸确保最终输出是 224×224。如果确实需要支持多种输入尺寸导出时把dynamic_axes里的高度和宽度也设为动态但这样会失去一些优化机会推理速度会下降。4.4 GPU 推理比 CPU 还慢现象装了onnxruntime-gpu但推理速度没有提升甚至更慢。原因模型太小数据传输开销超过了计算节省的时间。垃圾分类用的 MobileNet 或 ResNet18 参数量少GPU 的 kernel launch 开销占比高。另外如果每次推理都新建 sessionGPU 初始化时间也会拖慢整体速度。解决先确认providers里CUDAExecutionProvider排在第一位。然后把 session 做成全局变量只初始化一次。如果模型确实小CPU 推理反而更划算。我实测过 MobileNetV3 在 CPU 上单张推理 15ms 左右GPU 上 8ms但加上数据传输和初始化批量处理时 GPU 才有明显优势。4.5 转 RKNN 或 NCNN 后精度下降现象ONNX 模型在 ONNX Runtime 上精度正常转成 RKNN 或 NCNN 后精度明显下降。原因量化。RKNN 和 NCNN 默认会做 int8 量化量化过程会引入误差。垃圾分类模型如果对某些类别区分度不高量化后容易混淆。解决先确认转换时是否开启了量化。如果开启了尝试关闭量化或者用混合量化只量化部分层。另外转换工具对 ONNX 的 opset 版本有要求RKNN 通常要求 opset 11 或 12NCNN 对 11 支持最好。转换前先用onnxsim简化模型去掉冗余算子能提高转换成功率。5. 把 ONNX 垃圾分类模型塞进边缘设备的几个技巧5.1 用 onnxsim 给模型“瘦身”导出后的 ONNX 文件里往往有很多冗余算子比如连续的 Reshape、Transpose或者常量折叠没做干净的节点。onnxsim可以自动简化计算图减小文件体积有时还能提升推理速度。pip install onnxsim onnxsim garbage_classifier.onnx garbage_classifier_sim.onnx执行后会输出简化前后的算子数量对比。我遇到过导出后 120 个算子简化后剩 80 个的情况文件从 45MB 降到 42MB。简化后的模型精度不变因为 onnxsim 只做等价变换。5.2 转 RKNN 的完整命令与参数如果你要把模型跑在瑞芯微的 NPU 上需要转成 RKNN 格式。转换工具是rknn-toolkit2在 x86 机器上运行。from rknn.api import RKNN rknn RKNN() rknn.config( mean_values[[0.485 * 255, 0.456 * 255, 0.406 * 255]], std_values[[0.229 * 255, 0.224 * 255, 0.225 * 255]], target_platformrk3588 ) rknn.load_onnx(modelgarbage_classifier_sim.onnx) rknn.build(do_quantizationTrue, datasetquant_dataset.txt) rknn.export_rknn(garbage_classifier.rknn)mean_values和std_values要和训练时的 Normalize 对应注意这里乘了 255因为 RKNN 内部处理的是 0 到 255 的像素值。do_quantizationTrue会做 int8 量化需要提供一个quant_dataset.txt里面是几十张训练图片的路径用于校准量化参数。量化后模型体积缩小到原来的四分之一左右推理速度提升明显但精度可能掉 1 到 3 个百分点。如果精度掉得太多把do_quantization设为False用 float16 推理精度基本无损速度也有提升。5.3 一个容易被忽略的细节类别顺序ONNX 模型输出的类别顺序取决于你训练时CLASS_NAMES的定义顺序。这个顺序在导出、推理、转 RKNN 的整个链路里必须保持一致。我见过有人在训练时用[可回收, 厨余, 有害, 其他]推理时写成了[厨余, 可回收, 有害, 其他]结果所有分类都错位了。建议把类别名称写在一个单独的 JSON 文件里训练和推理都从这个文件读。这样改的时候只改一处不会漏。{ class_names: [可回收物, 厨余垃圾, 有害垃圾, 其他垃圾, 纸张, 塑料] }推理代码里用json.load读进来不要硬编码在 Python 文件里。这个习惯能帮你省下不少排查时间。5.4 验证边缘设备推理结果的方法模型转到 RKNN 或 NCNN 之后怎么确认结果是对的最直接的方法是在 PC 上用 ONNX Runtime 跑一遍记录每张测试图的输出向量然后在边缘设备上跑同样的图比较两个输出向量的余弦相似度。相似度在 0.99 以上算正常低于 0.95 说明转换过程有问题。如果没有条件在边缘设备上逐张对比至少要在转换后用 RKNN 的模拟器跑一遍验证集看准确率掉了多少。掉 1 到 2 个点可以接受掉 5 个点以上就要检查量化配置和预处理参数了。我自己做垃圾分类项目时习惯在导出 ONNX 之后先不急着转其他格式而是用 ONNX Runtime 跑完整验证集确认精度和 PyTorch 一致。这一步过了后面转 RKNN 或者 TensorRT 才有意义。如果 ONNX 这关就没过后面全是白费功夫。希望帮到你。本文还有配套的精品资源点击获取