垃圾分类的图像识别,难点从来不在模型结构本身。卷积神经网络做垃圾图像分类这件事,十年前就有人做,今天开源代码一抓一大把,但真正把它做成一套能跑在分拣设备上、误判率可控、后续能持续迭代的系统,中间隔着数据、评估口径、部署一致性这三道坎。我前后接触过几个类似的需求场景,有的做社区回收设备,有的做后端分拣线的辅助识别,最后发现模型准不准只是表象,真正决定成败的是你有没有把“一张照片从摄像头进来、到分类结果出去”这条链路当成一个工程问题来对待。这篇文章面向的是想动手复现一套卷积神经网络垃圾图像分类系统的开发者,也适合已经跑通了demo、但发现现场效果和验证集差一大截的同学。下面我按照自己实际搭这套东西的顺序,把设计思路、数据环节、训练细节、完整实操流程、踩过的坑和排查手段完整讲一遍,代码和参数都能直接抄。
1. 系统整体设计与技术路线拆解
1.1 为什么这类任务优先选卷积神经网络而不是传统特征工程
垃圾图像分类的本质是纹理加形状加材质的联合判别。玻璃瓶和塑料瓶在轮廓上高度相似,金属罐和纸盒在压扁状态下也经常撞脸,真正区分它们的是反光特性、透光性、边缘锐利度、表面纹理这些细节。传统方案走的是人工设计特征的路子,比如提取颜色直方图、LBP纹理、HOG边缘,再喂给SVM或者随机森林。我最早试过这条路,在实验室的干净背景数据集上能到七成多,一到真实场景就崩,原因很简单:光照一变,颜色直方图整体偏移;拍摄角度一变,HOG的方向分布全乱。手工特征对成像条件是强依赖的,而垃圾图像的成像条件恰恰是最不可控的那一环。
卷积神经网络的价值在于它把特征提取和分类器联合起来优化,卷积核在训练中自动学到的就是“什么样的局部模式对区分这几类有用”。浅层学到边缘和色块,深层学到材质纹理和部件组合,这套层级特征对光照、角度、轻微遮挡都有更好的容忍度。更重要的是迁移学习的存在,让这个任务不再需要从零训练。在千万级图像上预训练过的骨干网络,已经把通用的视觉特征学得差不多了,你只需要用几千张垃圾图像去微调最后的判别部分,效果就能甩开手工特征一大截。这也是我后来所有类似项目都直接用预训练骨干的原因,省下来的时间和算力远比多调几个超参数有价值。
需要说清楚的是,卷积神经网络不是万能药。它的优势建立在“训练数据分布和线上数据分布接近”这个前提上。如果你的训练集全是白底摆拍,线上全是复杂背景带遮挡,再强的网络也救不回来。这一点我在第 2 章和第 5 章会反复强调,因为它是这类项目翻车最常见的原因。
1.2 端到端系统的四层架构与选型逻辑
我把整套系统拆成四层:数据层、训练层、评估层、推理层。数据层负责采集、清洗、标注、划分,训练层负责增强、加载、前向反向、保存检查点,评估层负责指标计算、混淆分析、错误样本回流,推理层负责预处理、模型前向、后处理与拒识。这么拆不是为了好看,是为了让每一层都能单独替换和单独验证。实际项目中你会发现,换一个骨干网络只需要动训练层的一行配置,换一个摄像头只需要动推理层的预处理参数,如果代码揉成一坨,每次改动都要回归测试全链路,效率极低。
选型上的几个关键决策我列一下。骨干网络用预训练权重加微调,不从零训练,理由前面说过。输入分辨率定在 224×224,这是预训练模型的默认尺度,改了要重新调归一化参数,收益不明显。分类头用全局平均池化加一层全连接,参数量小、过拟合风险低。损失函数默认交叉熵,类别不均衡时加类别权重或换 Focal Loss。优化器用 AdamW 起步,收敛快、对学习率不敏感,等模型结构定下来要冲最后一两个点精度时再换 SGD 加余弦退火。这套组合不是唯一解,但对绝大多数中小规模垃圾图像数据集来说,是性价比最高的一条路。
还有一点选型上的取舍值得单独说:要不要上目标检测。很多人的第一反应是“垃圾可能一张图里有好几个,得先检测再分类”。我的建议是先做分类,把单物体分类的准确率和鲁棒性打扎实。检测模型标注成本高一个量级,而且要框得准,垃圾这种形变严重的物体框标注一致性很差。等分类模型跑通了,如果业务确实需要多物体场景,再在前面挂一个轻量检测器,或者用滑窗加分类头的粗暴方案过渡,都比一上来就啃检测要务实。
1.3 先定评估口径,再谈模型好坏
这是我最想强调的一条经验。很多人上来就盯 Top-1 准确率,跑出个 92% 就觉得很满意,结果上线发现有害垃圾被误判成其他垃圾的比例高得离谱。问题出在评估口径和业务目标脱节。垃圾分拣这件事,不同类别的误判成本完全不对称:把塑料瓶判成玻璃瓶,损失是分拣线上多一道人工复核;把电池判成其他垃圾,损失可能是整批处理的合规问题。所以评估指标必须分层看。
我通常要求同时看四个东西:整体 Top-1 准确率、每一类的召回率、宏平均 F1、以及混淆矩阵。宏平均 F1 能防止模型在多数类上刷分掩盖少数类的糟糕表现。混淆矩阵能直接告诉你哪两类在互相打架,这是后续做针对性优化的入口。另外我会单独统计一个“高置信度错误率”,也就是模型给出置信度大于 0.9 但判错的样本占比,这个数字直接决定了你能不能放心地把置信度当阈值用。
具体做法是维护一份固定的测试集,从第一版模型开始就不动它,所有模型都在同一份测试集上比。测试集的划分方式在第 2 章会详细讲,这里只提一句:随机划分图片是错的,必须按物体或按拍摄批次划分,否则测试集里的物体在训练集里出现过,指标虚高,上线必崩。
2. 数据集构建:决定上限的关键环节
2.1 类别体系怎么定:粗分类还是材质细分类
类别体系的设计直接决定了任务难度和最终可用性。常见的有两套思路。一套是按处置去向分,比如可回收物、厨余垃圾、有害垃圾、其他垃圾这四类。另一套是按材质分,比如玻璃、纸张、纸板、塑料、金属、其他这六类。两套体系各有适用场景,但如果你的目标是做一个图像分类系统,我强烈建议先做材质细分类,再把材质映射到处置类别。
原因在于类内差异。四分类体系里,“可回收物”这一类包含了塑料瓶、玻璃瓶、易拉罐、纸箱、旧书,这些物体在视觉上几乎没有共同点,模型要学的是一个极其宽泛的概念,效果必然打折。而材质分类体系里每一类的视觉特征相对集中,模型学起来容易得多,准确率也更容易做上去。映射关系是纯规则,玻璃、纸张、纸板、塑料、金属映射到可回收物,厨余和其他单独处理,写个字典就完事,零成本。
如果业务场景明确只需要四分类,那就在材质模型输出之后加一层映射,而不是直接训练四分类模型。这样做还有一个好处:后续业务调整分类口径时,你不需要重新标注和重训,改映射表就行。我在一个项目里遇到过处置规则中途调整的情况,因为底层是材质模型,改配置十分钟搞定,如果当初直接训的四分类,等于推倒重来。
类别确定之后还要考虑一个问题:要不要加一个“无法识别”类。我的建议是加。真实场景里总会有模糊、遮挡、非垃圾物体的输入,强行归到六类里任何一类都会产生错误。加一个“其他/不确定”类,用低置信度样本和人工构造的困难样本填充,模型学到的决策边界会更保守,上线后的可控性更好。
2.2 采集与标注的实操细节
公开数据集可以用来起步,但几乎不可能直接满足真实场景。公开数据集的拍摄条件往往比较理想,白底、正面、光线均匀,而你的摄像头装在设备上,角度是俯视或者斜视,背景是传送带或者桶口,光照受现场影响很大。我的做法是先用公开数据集把整条链路跑通、把代码调对,然后立刻转入自采数据。自采数据量不需要很大,每一类几百到一千张,但必须覆盖真实的拍摄条件。
采集时要注意几个点。第一,同一个物体要在不同角度、不同距离、不同光照下多拍几张,但不要把这几张分散到训练集和测试集里,这是数据泄漏的根源,下一节详细讲。第二,要刻意采集困难样本:反光的塑料瓶、压扁的纸盒、沾了污渍的玻璃、叠在一起的易拉罐。这些样本在真实分布里占比不低,但在随手拍的数据集里几乎见不到。第三,要采集负样本,也就是非垃圾物体、空传送带、局部遮挡的画面,用来支撑“其他”类。
标注环节最容易被低估。单人或单人连续标注会产生系统性偏差,比如同一个人对“纸板”和“纸张”的边界判断前后不一致。我的做法是每类抽 10% 做双人标注,计算一致性,低于 90% 就把边界规则写清楚,重新对齐口径。标注工具用什么都行,关键是导出格式统一,我习惯用“每类一个文件夹”的目录结构,文件名保留原始编号,方便后续追溯。
还有一个实操技巧:给每张图记录元信息,包括拍摄设备、拍摄时间、拍摄位置编号、光照条件标签。这些信息在做数据划分和错误分析时非常有用。比如你发现某台设备拍的图错误率特别高,可能是镜头脏了或者白平衡漂了;发现某个时间段的数据普遍差,可能是那批光照条件特殊。没有元信息,你只能看到一堆图片,无从下手。
2.3 数据泄漏:最容易被忽略的致命问题
数据泄漏这个坑,我见过太多项目栽在上面,而且栽了还不知道。典型场景是这样:你有一批垃圾图像,每个物体拍了五张不同角度的照片,然后你随手用随机划分把 70% 放训练、15% 放验证、15% 放测试。结果同一个塑料瓶的三张照片在训练集里,两张在测试集里。模型在训练时已经见过这个瓶子的所有细节,测试时再认一遍当然准,指标能到 98%,上线之后掉到 70%。这时候你会怀疑模型、怀疑数据增强、怀疑学习率,唯独想不到是划分方式的问题。
正确的做法是按物体 ID 或者按拍摄批次划分。每张图在采集时记录它属于哪个物体、哪一批次,划分时以物体为最小单位,同一物体的所有图片只能出现在一个集合里。拍摄批次同理,同一时间段、同一设备、同一光照条件下拍的一批图,最好整体划到同一个集合,否则模型会学到批次相关的捷径特征。这样划分出来的测试集指标会难看一些,但那个数字才是可信的。
判断有没有泄漏,有个简单的自检方法:看训练集和测试集的指标差距。如果训练集 99%、测试集 98%,在小数据集上这几乎一定是泄漏。正常的差距,几千张图、六分类、用迁移学习,训练集 95%、测试集 88% 到 92% 是合理区间。差距过大先查泄漏,再查过拟合。
2.4 数据增强与类别不均衡处理
增强策略的设计要贴着任务特点来。通用的随机裁剪、随机翻转、颜色抖动是标配,但强度要控制。垃圾图像对颜色是敏感的,塑料的本色、纸张的泛黄程度都是有判别价值的线索,颜色抖动的 hue 幅度我一般控制在 0.05 到 0.1,饱和度和亮度可以稍大一些。翻转要注意语义,水平翻转基本安全,垂直翻转对垃圾图像也安全,因为垃圾没有固定的上下方向,这一点和行人识别、文字识别不一样,可以放心用。
更强的增强手段包括 RandAugment、MixUp、CutMix。RandAugment 能在不写一堆策略的情况下自动组合增强操作,实测在小数据集上提升明显。MixUp 和 CutMix 对分类任务有效,能改善决策边界的平滑性,但要注意它们的标签是软标签,和标签平滑叠加时权重别加太重,否则收敛变慢。我的经验是 CutMix 更适合垃圾这种物体占画面主体的任务,MixUp 在背景复杂的场景里容易产生不自然的混合纹理。
类别不均衡是这类数据集的家常便饭。金属往往比纸张少很多,有害垃圾更是稀缺。处理手段有三种,我一般组合使用:重采样、类别权重、Focal Loss。重采样用 WeightedRandomSampler,按类别频率的倒数给样本加权,让每个 batch 里各类大致均衡,简单有效。类别权重直接传给交叉熵的 weight 参数,权重取频率倒数的归一化值。Focal Loss 适合难样本多的情况,公式是 FL = -α(1-p)^γ·log(p),γ 取 2 左右,能把大量易分样本的损失压下去,让模型聚焦在难例上。三种手段别一起上,容易过犹不及,先用重采样加类别权重,效果不够再换 Focal Loss。
3. 模型选型与训练策略的核心细节
3.1 骨干网络对比:ResNet、MobileNet、EfficientNet怎么选
骨干网络的选择要同时看精度、速度、部署条件三个维度。我整理了一张对比表,数值是公开基准上的参考值,实际项目里要用自己的数据实测。
| 骨干网络 | 参数量 | 输入尺寸 | 参考 Top-1 | CPU 单张延迟量级 | 适用场景 |
|---|---|---|---|---|---|
| ResNet50 | 约 25.6M | 224×224 | 76% 上下 | 百毫秒级 | 服务器端、精度优先 |
| MobileNetV3-Large | 约 5.4M | 224×224 | 75% 上下 | 几十毫秒级 | 边缘设备、速度优先 |
| EfficientNet-B0 | 约 5.3M | 224×224 | 77% 上下 | 几十毫秒级 | 精度与体积均衡 |
| ViT-B/16 | 约 86M | 224×224 | 81% 上下 | 数百毫秒级 | 数据量大、算力充足 |
表格里的数字只是起点,别当成结论。垃圾图像数据集通常只有几千张,在这种规模下,大模型的优势发挥不出来,反而更容易过拟合。ViT 这类结构对数据量的要求很高,没有几万张以上、没有强增强,微调效果往往不如 ResNet。我在一个小规模数据集上做过对比,EfficientNet-B0 和 MobileNetV3-Large 的最终精度几乎打平,但 MobileNetV3 的推理速度快一倍,最后选了后者。
选型的判断逻辑是这样的:先确认部署环境的算力预算。如果是服务器端批量处理,ResNet50 或 EfficientNet-B2 都行,精度优先。如果是嵌入式设备或者移动端,MobileNetV3 系列或者 EfficientNet-B0 起步,先看速度能不能满足帧率要求,再看精度。如果两边都要兼顾,可以做知识蒸馏,用大模型教小模型,但这个属于第二阶段优化,第一版别碰。
3.2 迁移学习的冻结策略与微调节奏
预训练模型的微调不是一上来就全部放开。我的标准流程分三段。第一段冻结整个骨干网络,只训练最后的分类头,学习率设大一点,比如 1e-3,跑 3 到 5 个 epoch。这一步的目的是让随机初始化的分类头先收敛到一个合理的位置,避免一开始就用大梯度把预训练好的卷积核冲乱。第二段解冻最后一个 stage,学习率降到 1e-4,再跑 5 到 10 个 epoch,让高层特征适应垃圾图像的纹理。第三段全部解冻,学习率再降一个量级到 1e-5,跑 10 到 15 个 epoch 收尾。
这个节奏背后的逻辑是分层学习率。浅层卷积学到的是边缘、颜色这类通用特征,几乎不需要改;深层卷积学到的是语义组合,需要大幅调整。如果全网络用同一个学习率,浅层要么学不动、要么被破坏,深层要么学得太慢、要么震荡。PyTorch 里可以给不同参数组设置不同学习率,实现起来很简单。
批归一化层要特别小心。BN 的 running mean 和 running var 是在预训练数据上统计的,微调时如果 batch size 太小,比如小于 16,统计量会很不准,反而伤害性能。这种时候有两个选择:冻结 BN 层的统计量不更新,只让可学习的缩放平移参数参与训练;或者把 BN 换成 GroupNorm。我一般在小 batch 场景下直接冻结 BN 统计量,改动最小,效果最直接。
3.3 损失函数、优化器与学习率调度
损失函数默认交叉熵,加标签平滑,smoothing 取 0.1。标签平滑的作用是防止模型过度自信,改善置信度的校准性,这对后面要用置信度做拒识很重要。不加标签平滑的模型,置信度分布会非常极端,几乎全是 0.99,阈值根本没法设。
优化器分阶段用。前面说的三段微调,前两段用 AdamW,weight decay 设 0.01 到 0.05,学习率按上面说的量级。最后一段如果追求极致精度,换成 SGD,momentum 0.9,学习率 0.01,配合余弦退火。AdamW 收敛快但最终精度有时略逊于调好的 SGD,这个差异在小数据集上不一定明显,看你有没有时间调。
学习率调度用带 warmup 的余弦退火。warmup 的作用是训练初期让学习率从很小的值线性升到设定值,避免一开始梯度爆炸。warmup 步数取总步数的 5%,或者固定 5 个 epoch,看哪个大用哪个。余弦退火从设定值平滑降到 1e-6,比阶梯下降更平滑,实测收敛曲线更好看。
批量大小和学习率是联动的。经验法则是线性缩放,batch 翻倍学习率翻倍。但这个法则在微调场景下不完全适用,因为预训练模型的容忍度更低。我的做法是 batch 设 32 到 64,学习率按上面的量级直接给,不硬套缩放规则,然后看前几个 epoch 的损失曲线,震荡就减半,下降太慢就加倍。
3.4 训练过程的监控指标与早停
训练过程中要盯的不只是 loss 和 accuracy。我一般记录这几个量:训练损失、验证损失、验证宏 F1、各类召回率、学习率、梯度范数。梯度范数很有用,如果它突然飙高,说明有异常样本或者学习率过大,配合梯度裁剪能防住大部分训练崩溃。
验证频率建议每个 epoch 跑一次完整验证,数据量大可以每隔几百步跑一次部分验证。验证时模型切 eval 模式,关掉 dropout 和 BN 的训练行为,这个低级错误每年都有人犯。
早停的判据用验证集宏 F1,连续 8 到 10 个 epoch 不提升就停,同时保存验证 F1 最高的那个检查点,而不是最后一个。最后一个 epoch 的模型往往已经轻微过拟合,用最高分检查点更稳妥。检查点保存时把优化器状态、当前 epoch、最佳指标一起存下来,方便中断续训。
还有个小技巧:训练结束后用测试集评估之前,先看看训练集和验证集的学习曲线。如果两条曲线贴得很近但都很高,说明模型容量不够或者任务太简单;如果训练高验证低且差距持续扩大,是过拟合,加增强、加正则、减容量;如果两条都低,是欠拟合,加容量、加训练轮数、查标签。这三种形态对应完全不同的处理方向,盲目调参效率极低。
4. 完整实操流程:从环境搭建到模型导出
4.1 环境与目录结构
环境用 Python 3.9 以上,PyTorch 2.x,torchvision 配套版本,显卡驱动和 CUDA 版本对齐。依赖不多,除了框架本身,装一下 numpy、Pillow、scikit-learn 用于指标计算,tqdm 用于进度显示就够了。不建议一上来就装一堆实验管理工具,第一版跑通比什么都重要。
目录结构我习惯这样组织,清晰且方便脚本引用:
trash_cls/ ├── data/ │ ├── train/ │ │ ├── glass/ │ │ ├── paper/ │ │ ├── cardboard/ │ │ ├── plastic/ │ │ ├── metal/ │ │ └── trash/ │ ├── val/ │ └── test/ ├── splits/ │ ├── train.csv │ ├── val.csv │ └── test.csv ├── src/ │ ├── dataset.py │ ├── model.py │ ├── train.py │ ├── evaluate.py │ └── export.py ├── configs/ │ └── base.yaml └── outputs/splits目录里放划分清单,每行是图片路径和标签,划分逻辑写死在脚本里,保证每次实验用的是同一份划分。这一点很重要,我见过有人每次跑实验都重新随机划分,结果两个模型的指标根本没法比。
4.2 数据管道与增强配置
数据管道是性能瓶颈的高发区。图片解码是 CPU 密集操作,如果 num_workers 设得太小,GPU 会一直等数据。一般设成 CPU 核数的 0.75 倍左右,配合 pin_memory 和 non_blocking 传输,能把 GPU 利用率拉满。
训练和验证的预处理必须分开写,而且要和推理端严格对齐。训练端用随机裁剪加增强,验证端用短边缩放加中心裁剪,这是标准做法:
import torch from torchvision import transforms IMAGENET_MEAN = [0.485, 0.456, 0.406] IMAGENET_STD = [0.229, 0.224, 0.225] train_tf = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.6, 1.0), ratio=(0.75, 1.33)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomVerticalFlip(p=0.2), transforms.RandAugment(num_ops=2, magnitude=7), transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.05), transforms.ToTensor(), transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), transforms.RandomErasing(p=0.25, scale=(0.02, 0.15)), ]) val_tf = transforms.Compose([ transforms.Resize(256, interpolation=transforms.InterpolationMode.BILINEAR), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ])几个细节解释一下。RandomResizedCrop 的 scale 下限定在 0.6,不要设太低,垃圾物体通常占画面主体,裁得太狠会裁掉判别性部分。RandAugment 的 magnitude 我设 7,比默认的 9 温和一些,因为垃圾图像的纹理和颜色本身就有判别价值,增强太猛反而破坏信息。RandomErasing 放在 ToTensor 和 Normalize 之后,这是 torchvision 的接口要求,顺序错了会报错。
类别不均衡用采样器处理:
import numpy as np from torch.utils.data import WeightedRandomSampler def build_sampler(labels, num_samples=None): labels = np.asarray(labels) classes, counts = np.unique(labels, return_counts=True) class_weight = {c: 1.0 / n for c, n in zip(classes, counts)} sample_weight = np.array([class_weight[l] for l in labels], dtype=np.float64) if num_samples is None: num_samples = len(labels) return WeightedRandomSampler( weights=torch.as_tensor(sample_weight, dtype=torch.double), num_samples=num_samples, replacement=True, )采样器只用在训练集上,验证集和测试集绝对不能用,否则评估结果失真。
4.3 训练循环核心代码
训练循环要处理混合精度、梯度裁剪、梯度累积这几件事。混合精度用 torch 自带的 autocast 和 GradScaler,能在几乎不掉精度的情况下省一半显存、提一截速度。梯度裁剪防梯度爆炸,阈值设 5.0 比较通用。
import torch from torch.cuda.amp import autocast, GradScaler def train_one_epoch(model, loader, criterion, optimizer, scheduler, device, scaler, clip=5.0): model.train() total_loss, correct, total = 0.0, 0, 0 for images, labels in loader: images = images.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) with autocast(enabled=(device.type == "cuda")): logits = model(images) loss = criterion(logits, labels) optimizer.zero_grad(set_to_none=True) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), clip) scaler.step(optimizer) scaler.update() scheduler.step() total_loss += loss.item() * labels.size(0) correct += (logits.argmax(dim=1) == labels).sum().item() total += labels.size(0) return total_loss / total, correct / total这里 scheduler 是按步更新的,也就是每个 batch 走一次,这对应的是 cosine 加 warmup 的按步调度。如果你用的是按 epoch 更新的调度器,把它挪到 epoch 循环外面,别搞混。
学习率调度用 lambda 实现 warmup 加余弦:
import math def build_scheduler(optimizer, warmup_steps, total_steps, base_lr, min_lr=1e-6): def lr_lambda(step): if step < warmup_steps: return step / max(1, warmup_steps) progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) return max(min_lr / base_lr, cosine) return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)warmup_steps 取总步数的 5%,total_steps 等于 epoch 数乘以每个 epoch 的 batch 数,算的时候别漏了梯度累积的倍数。
4.4 显存与批量大小的估算方法
显存不够是新手最常见的报错。粗算的方法是:模型参数占一份,梯度占一份,优化器状态占一到两份,激活值占大头。以 ResNet50 为例,参数量约 25.6M,FP32 下权重约 100MB,梯度 100MB,AdamW 的两个状态各 100MB,加起来约 400MB。激活值跟 batch size 和分辨率成正比,224×224 输入、batch 64 的情况下,激活值大概在 3GB 到 4GB 之间。总计不到 5GB,16GB 的卡跑 batch 64 绰绰有余,开混合精度还能再省一截。
估算公式可以写成:激活显存约等于 batch_size 乘以单张图的激活量。单张 224×224 的 ResNet50 前向激活量大约 50MB 到 60MB,batch 64 就是 3.2GB 到 3.8GB,和实测吻合。这个数字随网络深度增加而增加,EfficientNet 系列的激活分布更不均匀,估算偏差会大一些。实际操作中不用算得太细,从 batch 32 起步,逐步往上翻倍,OOM 了就退回去,这是最快的办法。
显存吃紧时的处理顺序:先开混合精度,再减 batch 配合梯度累积,再考虑降分辨率,最后才换小模型。梯度累积的写法是每累积 N 个 batch 才 step 一次优化器,等效于把 batch 放大 N 倍,但要注意 BN 的统计量还是按实际 batch 算的,所以累积对 BN 没有帮助,这是它的局限。
4.5 模型导出与推理一致性校验
训练完之后要把模型导出成通用格式,方便部署端加载。导出本身很简单,难的是保证导出前后的输出一致。不一致的来源几乎全是预处理差异,这个问题在第 5 章会详细讲,这里先说导出和校验的流程。
import torch model.eval() dummy = torch.randn(1, 3, 224, 224, device=device) torch.onnx.export( model, dummy, "outputs/trash_cls.onnx", input_names=["input"], output_names=["logits"], opset_version=17, dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}}, do_constant_folding=True, )导出后必须做数值校验。取一批测试图片,分别在 PyTorch 和导出模型上跑前向,比较输出的最大绝对误差。正常情况下误差应该在 1e-4 量级以内,超过 1e-2 就说明某一层导出有问题,通常是动态 shape 或者算子不支持导致的。
校验代码大致这样:
import numpy as np import onnxruntime as ort session = ort.InferenceSession("outputs/trash_cls.onnx", providers=["CPUExecutionProvider"]) with torch.no_grad(): for images, _ in test_loader: images = images.to(device) ref = model(images).cpu().numpy() ort_out = session.run(None, {"input": images.cpu().numpy()})[0] diff = np.abs(ref - ort_out).max() print(f"max abs diff: {diff:.6f}") break这个校验我建议写进 CI,每次改模型或者改预处理都跑一遍,能挡住绝大多数部署事故。
5. 常见问题与排查技巧实录
5.1 训练损失不降或者反复震荡
损失不降,第一件事不是调学习率,是过拟合一个小 batch。取 16 张图,关掉所有增强和正则,训练几十步,看损失能不能降到接近 0。如果能,说明模型和代码没问题,问题在数据或超参;如果不能,说明代码有 bug,常见的有标签和图片错位、数据加载器返回的顺序乱了、损失函数用错、模型输出没接对。
标签错位是最隐蔽的一种。我遇到过一次,数据加载时用了两个独立的 shuffle,图片和标签各洗各的,模型学了两天损失一直在 1.7 附近打转。排查方法是打印前几个 batch 的图片和标签,肉眼看对不对得上,五分钟能定位。
损失震荡通常是学习率过大或者 batch 太小。先把学习率降一个量级看曲线是否变平滑,如果没有改善,检查是不是 batch 太小导致 BN 统计量波动。另外检查一下数据里有没有极端异常样本,比如全黑图、全白图、尺寸异常图,这些会制造巨大的梯度。
5.2 验证集指标很高但现场效果拉胯
这是最经典的域偏移问题。验证集和训练集同分布,现场数据和训练数据不同分布,指标自然对不上。具体表现可能是光照不同、背景不同、拍摄角度不同、设备色彩风格不同。排查的第一步是把现场错分的样本拿回来,和训练集放一起看,找出分布差异的具体维度。
背景偏差是最常见的一种。模型可能根本没在看垃圾,而是在看背景。比如玻璃瓶的训练图大多拍在浅色桌面上,模型学到的是“浅色背景就是玻璃”。验证集也是同样的背景,所以指标很好,现场背景一换就全错。验证方法是做 Grad-CAM 可视化,看模型的注意力落在哪。如果热力图集中在背景区域,基本可以确诊。
解决办法有三个层次。数据层面,采集带现场背景的数据,做背景替换增强,把前景物体抠出来贴到随机背景上。训练层面,加 CutMix、加更强的颜色和几何增强。评估层面,专门构造一个背景多样的测试集,用它来评估模型的真实鲁棒性。三个层次一起上效果最好,但数据层面的补充是根本,增强只能缓解不能根治。
5.3 两类之间混淆严重怎么定位
混淆矩阵是定位工具。假设你发现玻璃和塑料的互相误判占到总错误的四成,先别急着调模型,把这两类的错分样本全部拉出来看。常见原因有三类:标注边界不清、视觉特征确实接近、样本量不均衡。
标注边界不清的情况很普遍。透明的塑料杯和玻璃杯,标注的人自己都犹豫。这种要把规则写死,比如按材质成分判,按敲击声判不现实,就按外观的典型特征判,并且在标注文档里配上示例图。规则统一之后重新标注一批,混淆会明显下降。
视觉特征接近的情况,靠增强和更大分辨率解决。224 的分辨率下,玻璃的细小气泡和塑料的模具纹路看不清,升到 320 或者 384 会有改善,代价是显存和推理时间上升。另一个手段是引入辅助信息,比如在分类头之外加一个辅助分支预测“是否透明”,多任务学习能逼着模型学更细的特征。
样本量不均衡的情况,前面说的重采样和类别权重就是为此设计的。先把每类的样本数打出来看,差一个量级就必须处理。
5.4 部署后端精度掉点
模型在本地评估 90%,部署到推理端变成 75%,这类问题的排查有固定顺序,我整理成一张表。
| 排查项 | 具体检查内容 | 常见错误 |
|---|---|---|
| 通道顺序 | 模型期望 RGB,送入是否 BGR | 摄像头默认 BGR,未转换 |
| 归一化参数 | 均值和标准差是否与训练一致 | 用了 0.5/0.5 而非 ImageNet 参数 |
| 缩放方式 | 是否与验证集一致 | 训练用短边缩放,部署用直接拉伸 |
| 插值算法 | 双线性还是最近邻 | 部署端默认最近邻,引入锯齿 |
| 数据类型 | 是否做了归一化到 0 到 1 | 送了 0 到 255 的原始值 |
| 量化误差 | INT8 校准集是否覆盖真实分布 | 校准集全是训练图,未覆盖现场光照 |
通道顺序和数据类型这两项占了掉点原因的七成以上。我见过一个案例,部署端把 BGR 当 RGB 送进去,模型把蓝色通道当红色通道理解,整体精度掉了十五个点。排查方法很简单,用同一张图在两边跑,逐层对比中间输出,第一层就对不上,问题一定在预处理。做一次这样的对齐,比读十篇文档都管用。
5.5 高频问题速查表
把上面这些整合成一张速查表,遇到问题时按顺序过一遍,能省掉大量试错时间。
| 现象 | 优先排查 | 处理手段 |
|---|---|---|
| 损失不降 | 小 batch 过拟合测试 | 查标签错位、查学习率量级 |
| 损失震荡 | 学习率、batch 大小 | 降学习率、加 warmup、增大 batch |
| 训练高验证低 | 过拟合或数据泄漏 | 按物体重新划分、加增强、加正则 |
| 两类混淆严重 | 看错分样本和标注规则 | 统一标注口径、升分辨率、加辅助任务 |
| 现场效果差 | 背景偏差、光照偏移 | 补现场数据、背景替换、强增强 |
| 部署掉点 | 预处理一致性 | 逐层对比、对齐归一化和缩放 |
| 显存溢出 | batch 和分辨率 | 混合精度、梯度累积、降分辨率 |
| 推理太慢 | 模型规模和输入尺寸 | 换轻量骨干、量化、降分辨率 |
6. 工程化落地与效果评估的几点经验
6.1 置信度阈值与拒识机制的设计
分类模型输出的是 softmax 概率,很多人直接拿最大值当置信度用,设个 0.6 的阈值,低于就拒识。这个做法有前提:模型的置信度得是校准过的。没校准的模型,置信度普遍虚高,一万个样本里可能九千九百个都是 0.99,阈值形同虚设。
校准的方法用温度缩放最简单。在验证集上拟合一个温度参数 T,推理时把 logits 除以 T 再算 softmax,能让置信度分布更接近真实准确率。实现上就是保存一个标量,成本极低。做完校准之后,你可以画一条可靠性曲线,横轴是置信度分桶,纵轴是该桶内的实际准确率,曲线贴近对角线说明校准良好。
拒识阈值怎么定,要看业务能接受多少人工复核量。我的做法是定两个阈值:高阈值以上直接输出,低阈值以下直接拒识,中间区间走二次确认。比如置信度大于 0.85 直接判,小于 0.4 直接拒,中间的交给人或者交给更复杂的模型。这样能把绝大部分简单样本快速处理掉,把算力和人力集中在困难样本上。这套机制在实际系统里的价值远大于把准确率再刷高一个点。
6.2 端侧部署的量化取舍
如果部署目标是嵌入式设备或者移动端,量化是绕不开的一步。INT8 量化通常能把模型体积压到四分之一,推理速度提升两到四倍。代价是精度会掉,掉多少取决于校准集的质量和模型对量化的敏感程度。
我的经验是,分类任务上 INT8 量化掉点控制在 1 个点以内是完全可以做到的,前提是校准集要有代表性。校准集不能只用训练集,必须包含现场采样的真实图片,覆盖不同光照、不同背景、不同设备。校准样本数量 500 到 1000 张足够,太少统计不准,太多收益递减。
对量化敏感的层通常是第一层卷积和最后的全连接层,可以把这两层保留成 FP16 或者 FP32,中间层用 INT8,混合精度量化能把掉点压到最小。另外 BN 层在量化时会被折叠进卷积,这个折叠过程要确保用的是推理模式的统计量,用训练模式的统计量会引入明显误差。
6.3 数据闭环与持续迭代的落地方式
第一版模型上线只是开始。真实场景的数据分布会漂移,新包装材料出现、季节变化导致光照不同、设备老化导致色彩偏移,模型的表现会缓慢下降。要维持效果,必须建立数据闭环。
闭环的具体做法是:推理端记录每一条预测的置信度和输入图片的缩略图,低置信度的样本自动回流到待标注池,定期人工标注后加入训练集,每月或者每季度做一次增量训练。增量训练不是从头训,是在上一版检查点的基础上继续微调,学习率调小,同时保留一部分旧数据防止灾难性遗忘。
版本管理要跟上。每次训练记录数据版本、配置、指标、模型文件,形成可追溯的链条。我见过团队因为没做版本管理,线上出问题想回滚,结果找不到上一版模型对应的训练数据,只能干着急。这套流程看起来繁琐,但它是模型能长期可用的基础,比任何一次调参带来的提升都重要。
我个人在实际操作中的体会是,这套系统里模型结构大概只占两成工作量,数据工程和部署一致性占八成。很多人把时间全花在换骨干网络、调学习率上,收益微乎其微;反而是把数据划分做对、把预处理对齐、把拒识机制加上,效果提升立竿见影。如果你正准备动手,我的建议是先花两天把数据管线和评估流程搭扎实,再花一天跑通第一版模型,剩下的时间全部投在错误分析和数据补充上。模型选型这件事,等你的测试集指标稳定在九成以上、混淆矩阵的错分集中在特定类别时,再回来调也不迟。