基于ShuffleNet的菠萝成熟度8档分类实战

简介面向计算机视觉初学者与轻量级网络研究者该实战项目基于ShuffleNet完成8种不同成熟阶段菠萝的图像分类覆盖未熟、半熟、成熟等类别。完整提供可直接运行的Python代码、标注图片数据集与训练好的权重文件训练集4808张、测试集806张目录按类别划分清晰。ShuffleNet参数量仅约一百万采用余弦学习率衰减训练50个epoch测试集最高精度达87%run_results目录中保存最佳权重、训练日志及loss/精度曲线便于复现与调优。资源共2000个文件以jpg图像为主另有4个py脚本、1个json配置及readme说明压缩包约201MB已有124人学习。预测阶段运行predict即可自动推理inference中全部图片并在左上角标注前三个概率最大类别如需训练自定义数据readme中已说明自动生成类别数量的流程。该资源适合快速上手轻量级分类任务也可作为算法对比与课程设计的可靠参考。1. 用轻量级CNN把菠萝成熟度分成8档为什么偏偏是ShuffleNet做农业智能分拣时菠萝成熟度分级是个容易翻车的任务。它不像猫狗分类那样类别差距大相邻两个成熟阶段之间颜色可能只差几个 RGB 值光线一变就可能判错。我接到的需求是 8 阶段分类从全绿到深褐按成熟度分批采收、分批上市。部署环境是一台几百块钱的边缘盒子不可能跑 VGG 或 ResNet-152。ShuffleNet 这种轻量级 CNN 就成了首选它在不牺牲太多精度的前提下把单位计算量的信息量拉满。这篇实战记录会把整个方案按流程拆开新手能照着跑通老手能看参数边界和坑。2. 数据准备把8种成熟度阶段变成模型能学的目录结构2.1 8个阶段怎么定先画出分级标准再拍图做分类任务最先碰到的不是模型而是“成熟度阶段”这个边界。用经验拍脑袋定 8 档模型学到的一定是噪声。我一般会先把标准定成能被任何人复述的颜色和面积描述再让产线上的同事确认。下面的表是我在菠萝产线上常用的分级基准你可以按自己的品种调整。阶段判定依据典型颜色描述1全果深绿无黄色区域深绿色果眼未变亮2黄色面积 20%出现零星黄晕绿多黄少顶部微黄3黄色面积 20% ~ 40%基部先转暗黄黄绿相间4黄色面积 40% ~ 60%颜色偏浅浅黄为主仍有斑驳绿5黄色面积 60% ~ 80%菠萝整体明亮明亮的黄6黄色面积 80%果眼周围偏橙橙黄光泽明显7出现褐色斑点褐斑面积 20%黄褐相间果皮开始发暗8褐斑连片果皮深褐可能有软化区深褐表面不均匀这个分级的关键是相邻阶段必须有肉眼可判断的差异否则标注员和算法都会被逼疯。第 2 和第 3 阶段之间看起来只是黄多黄少的差别但如果你把标准细化成“基部转黄”这样一个位置描述标注一致性会明显提高。成熟度分类不是检测不需要坐标框只需要给每张图片一个标签但标签的质量决定了整个项目的上限。我还会给每个阶段拍一段包含不同角度的视频再抽帧而不是只拍几张正脸。菠萝的果眼和冠芽在不同角度下颜色变化很大只用标准正脸会让模型在产线上遇到俯视或侧视时崩溃。2.2 整理数据到ImageFolder代码与参数不管数据来自手机还是工业相机最终我都会把它整理成 PyTorch ImageFolder 的标准结构train/阶段名/图片val/阶段名/图片。阶段名用stage1到stage8这种前缀类名不要用中文否则要额外处理编码问题。我这里有一个现成的整理脚本import os import shutil from sklearn.model_selection import train_test_split raw_dir raw_pineapple # 原始图片目录文件名为 stage1_001.jpg 这种格式 train_dir data/train val_dir data/val stage_names [fstage{i} for i in range(1, 9)] for split_dir in [train_dir, val_dir]: for stage in stage_names: os.makedirs(os.path.join(split_dir, stage), exist_okTrue) # 按文件前缀归类并按 8:2 切分训练集和验证集 for stage in stage_names: files [f for f in os.listdir(raw_dir) if f.startswith(stage)] train_files, val_files train_test_split(files, test_size0.2, random_state42) for f in train_files: shutil.copy(os.path.join(raw_dir, f), os.path.join(train_dir, stage, f)) for f in val_files: shutil.copy(os.path.join(raw_dir, f), os.path.join(val_dir, stage, f))逻辑说明这里先把同一个阶段的图片全部读出来再按 8:2 切分。random_state42保证可复现。实际项目中我会先按拍摄批次分组再切分而不是直接随机切分。比如同一颗菠萝在第 1 天和第 4 天被拍摄这两张图应该放到同一个集合里。如果随机切分模型其实已经见过同一颗菠萝的另一成熟度版本验证集准确率就会虚高。整理完目录后先统计每个类的样本数量from collections import Counter from torchvision.datasets import ImageFolder train_ds ImageFolder(data/train) counts Counter([sample[1] for sample in train_ds.samples]) for i, name in enumerate(train_ds.classes): print(name, counts[i])这段代码能让你在训练前看到每个阶段的样本分布。如果你发现某个阶段只有另外一类的三分之一那就需要补拍或者做重采样不要直接开训。ImageFolder 默认按文件夹名称排序stage1到stage8的顺序是固定的后面计算混淆矩阵时不会乱。我这里还会加一步模糊图清洗。菠萝表面纹理多工业相机偶尔对焦失败模糊图会被模型当成一类噪声。用拉普拉斯方差做粗筛很实用import cv2 def is_blurry(path, threshold80.0): img cv2.imread(path, cv2.IMREAD_GRAYSCALE) if img is None: return True laplacian_var cv2.Laplacian(img, cv2.CV_64F).var() return laplacian_var thresholdthreshold80 是一个经验初始值不同相机的噪声水平不一样可以先随机看 20 张被筛掉的图再调阈值。宁可保留一些轻微模糊图也不要误删正常图因为产线上本身会有轻微振动。2.3 预处理与数据增强不让模型学会光线而学会颜色菠萝成熟度分类的特征主要是颜色分布其次才是果眼纹理。所以预处理的重点是颜色归一化和光照增强。常见做法是先把图缩放到 256x256再随机裁剪到 224x224验证集则不裁剪直接缩放到 224。这样可以保留一定的平移不变性又不会像直接变形那样扭曲果眼比例。from torchvision import transforms train_transforms transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.05), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])参数说明RandomCrop和Resize组合训练验证集直接Resize避免验证时因裁剪位置不同造成随机波动。ColorJitter的 brightness 和 saturation 对菠萝尤其重要因为产线上的白炽灯和自然光混在一起色温经常变化hue 只给了 0.05太小没有意义太大会让绿色偏成蓝色反而学坏。这里有一个容易踩的坑归一化用 ImageNet 的 mean/std 还是自己算的如果你加载了预训练权重就一定用 ImageNet 的 mean/std如果从头训练自己算也可以但我建议一律用 ImageNet 的这样以后换模型、加载预训练权重都不需要再改数据流。不要小看这些数字我见过有人用mean0.5, std0.5结果模型前几层的梯度波动非常大学会黄色区域的速度慢了一半。数据增强不要一开始就上 CutMix 或 Mixup。对成熟度分类颜色是强信号过度混合会让模型学到“绿黄之间的中间色”反而把相邻阶段抹平。先用基础的裁剪和颜色扰动如果过拟合再加更强的增强。3. 构建ShuffleNet分类网络通道混洗原理与PyTorch实现3.1 为什么ShuffleNet适合这种成熟度分类ShuffleNet 是一类专门为移动和边缘设备设计的轻量级 CNN。它的核心思路是用分组卷积减少计算量再用通道混洗让不同分组的信息互相流通。和 MobileNet 的深度可分离卷积相比ShuffleNet 在同样的计算量下精度略高一些尤其在没有专用 GPU 的盒子上这个差距会直接体现在帧率上。菠萝成熟度分类的场景有一个特点图像里没有复杂的空间结构不需要非常大的感受野但颜色过渡和果眼纹理又需要跨通道的配合。ShuffleNet 的通道混洗正好在这种尺度上做信息交换既不奢侈也不浪费。所以我没有选 ResNet18 或者 MobileNetV3而是用 ShuffleNetV2 作为主力模型。3.2 通道混洗核心模块从代码理解设计意图先看最核心的通道混洗实现。ShuffleNet 的分组卷积会让每个分组只处理自己的通道如果不混洗下一层的某个分组就看不到其他分组的特征模型表达能力会下降。通道混洗就是把这个过程变成一次 reshape 加 permuteimport torch def channel_shuffle(x, groups): batch, channels, height, width x.shape channels_per_group channels // groups # 先把通道拆成分组数再交换分组维度和通道维度 x x.reshape(batch, groups, channels_per_group, height, width) x x.permute(0, 2, 1, 3, 4).contiguous() return x.reshape(batch, channels, height, width)逻辑说明假设输入有 4 个通道分成 2 组reshape 成(batch, 2, 2, h, w)permute 后变成(batch, 2, 2, h, w)其实交换的是分组和组内通道两个维度最后展平回(batch, 4, h, w)。这样第一个分组里原本只有通道 0、1混洗之后会拿到通道 0、2 的数据。实际用的时候我不会手写整个 ShuffleNet因为 torchvision 的实现已经过充分验证。但理解这个模块对后面排查 ONNX 导出问题很重要通道混洗在推理引擎里容易生成大量 transpose 算子如果你的推理框架优化得不好速度会反而不如 MobileNet。3.3 用预训练ShuffleNetV2完成8类分类直接加载预训练模型替换最后的全连接层是最高效的路径。ShuffleNetV2 在 torchvision 里有不同宽度系数我用得最多的是x1_0如果边缘盒子性能太弱再降成x0_5。import torch.nn as nn from torchvision.models import shufflenet_v2_x1_0, ShuffleNet_V2_X1_0_Weights model shufflenet_v2_x1_0(weightsShuffleNet_V2_X1_0_Weights.DEFAULT) model.fc nn.Linear(model.fc.in_features, 8) print(model)逻辑说明shufflenet_v2_x1_0的默认分类头是 1000 类我们把fc替换成输出 8 维的新全连接层。in_features会动态取原来的输入维度不需要硬编码。这样保留前面所有卷积层的预训练权重只有最后的分类头是随机初始化的。如果你有多个 GPU记得先包DataParallel再替换否则参数名里会多了module.前缀后面保存和加载容易出错。如果你担心数据量少导致过拟合可以冻结一部分参数只训练后半段。先用一个小代码打印所有层名再决定冻结哪些层for name, param in model.named_parameters(): print(name, param.shape)不同 torchvision 版本的层名可能略有不同不要凭记忆写。常见做法是只让最后的fc和倒数第二个卷积模块可训练。冻结操作for name, param in model.named_parameters(): if fc not in name: param.requires_grad False这样冻结了全部卷积层只更新fc。如果你的数据只有几百张这种微调方式更稳如果数据量几千张可以只冻结前一半让后半段适应菠萝色彩分布。宽度系数怎么选x1_0的计算量大约是x0_5的四倍。菠萝成熟度分类的颜色特征并不需要特别细的纹理x0_5在大多数场景下就能达到 90% 以上的准确率。我通常会先跑一版x0_5快速验证数据质量再上x1_0做最终模型。不要在第一次训练就直接调最强模型浪费时间还不容易定位问题。4. 训练配置与调参让模型在8类上收敛4.1 损失函数与类别权重8 类成熟度分类还是标准的图像分类问题损失函数第一选择是交叉熵。但产线数据经常不均衡阶段 1 只拍几天阶段 4 能拍很久样本差距自然就出来了。当某个类的样本量明显偏少时我会用类别权重把损失拉平。from collections import Counter import torch import torch.nn as nn train_ds ImageFolder(data/train) counts Counter([sample[1] for sample in train_ds.samples]) total sum(counts.values()) weights torch.tensor([total / counts[i] for i in range(8)], dtypetorch.float) # 限制最大权重防止少数类被放大过头 weights torch.clamp(weights, max10.0) criterion nn.CrossEntropyLoss(weightweights.to(device))逻辑说明这里把每个类的样本数倒数乘以总样本数得到初步权重再用clamp限制最大权重为 10。如果某个类只有 20 张而多数类有 2000 张不限制权重的话损失会被放大 100 倍训练会震荡。加权交叉熵的本质是牺牲多数类的一些准确率去换少数类的召回所以如果前后两次训练权重差异太大记得回看混淆矩阵不要只看总体准确率。4.2 优化器、学习率与调度用预训练权重微调时我优先选 SGD 加 momentum而不是 Adam。SGD 的收敛在分类任务上更稳尤其是批次比较小的时候Adam 容易走到一个尖锐的极小值。学习率一般从 0.01 这个量级开始batch size 32如果 batch size 翻倍学习率也翻倍这是经验规则。权重衰减 5e-4 对卷积网络足够。import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4) scheduler CosineAnnealingLR(optimizer, T_max30)参数说明T_max30表示学习率会在 30 个 epoch 内从 0.01 按余弦曲线降到接近 0。这样做比固定学习率更好后期小步长能精修权重。如果训练数超过 30可以重新设置T_max或者用CosineAnnealingWarmRestarts做周期重启但我建议先跑一个完整周期看趋势。更稳妥的做法是加前几个 epoch 的 warmup。预训练权重已经在大数据集上收敛一上来就用 0.01 可能会让前几层剧烈抖动。常见做法是让学习率在头 5 个 epoch 从 0.001 线性升到 0.01def lr_lambda(epoch): if epoch 5: return 0.1 0.9 * epoch / 5 return 0.1 0.9 * 0.5 * (1 __import__(math).cos((epoch - 5) * 3.14159 / 25)) scheduler optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)这个函数前 5 个 epoch 做线性升压后面模拟余弦。注意这里的余弦周期和真实 epoch 数要匹配否则最后学习率降不到底。如果你不需要太精细直接用CosineAnnealingLR加一个LinearLR(optimizer, start_factor0.1, total_iters5)组合更省事。4.3 训练循环与检查点管理训练循环不需要花哨关键是保存最好的验证集模型而不是最后一个 epoch。我习惯用准确率作为早停指标因为 8 类均衡的情况下准确率足够直接。import os best_acc 0.0 patience 5 bad_epochs 0 os.makedirs(checkpoints, exist_okTrue) for epoch in range(30): model.train() train_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() loss criterion(model(images), labels) loss.backward() optimizer.step() train_loss loss.item() * images.size(0) model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) preds model(images).argmax(dim1) correct (preds labels).sum().item() total labels.size(0) acc correct / total print(fepoch {epoch 1}: loss {train_loss / len(train_ds):.4f}, acc {acc:.4f}) if acc best_acc: best_acc acc torch.save(model.state_dict(), checkpoints/best_shufflenet.pth) bad_epochs 0 else: bad_epochs 1 if bad_epochs patience: break逻辑说明训练损失按样本数加权平均验证集只看准确率。best_acc更新时保存state_dict而不是整个模型对象这样以后加载和导出 ONNX 都方便。patience5的作用是如果连续 5 个 epoch 没有刷新最好成绩就提前停止给后面的数据分析和调参留时间。训练时要注意验证集数据在 loader 里不能 shuffle我这里已经保持了shuffleFalse。还有一个小细节如果你开启了pin_memory在to(device)之前先用non_blockingTrue会更好但数据量不大时差别很小。混合精度训练我也经常开因为你不想为 8 类分类等一个晚上scaler torch.cuda.amp.GradScaler() ... with torch.cuda.amp.autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度能省一半显存但对 SGD 的动量参数和权重衰减没有影响可以放心用。如果你的 GPU 是旧架构不支持 AMP直接跳过这一段不影响训练结果。5. 菠萝成熟度分类避坑记录5个我踩过的坑5.1 模型把所有菠萝都判成阶段4现象训练 loss 稳中有降但验证集准确率卡在 30% 左右打印预测结果发现九成样本都判成了阶段 4。原因阶段 4 恰好是所有类别里样本最多的一类。交叉熵在类别不均衡时会把决策边界推到多数类一侧因为把所有样本都判成多数类整体 loss 依然不高。另一个原因是阶段 3、4、5 的颜色本身比较接近模型没有足够的梯度去细分边界。解决我先用混淆矩阵确认每个类的召回率只统计对角线就能看到阶段 1 和 8 的召回可能为 0。然后给CrossEntropyLoss加类别权重同时把阶段 4 的增强强度调高一点比如多做一些ColorJitter让模型不能只靠一个大类覆盖所有样本。加了权重之后整体准确率可能下降 1 到 2 个点但从产线上看阶段 1 被漏掉比阶段 4 被误判更严重。5.2 相邻阶段被严重混淆现象混淆矩阵里跨两个阶段以上的错判很少但阶段 2 和 3、5 和 6 之间的误差非常高几乎有一半被互相交换。原因相邻阶段的分界线本身就是人为定的颜色变化是连续区间没有硬性边界。标注员在标阶段 2 和阶段 3 时也可能意见不一致数据本身带了很大的标注噪声。加上自然光下不同的受光面会让同一个菠萝看起来像两个阶段。解决我重新收了一批数据让两个标注员独立标同一个菠萝只保留两人共识的样本对争议多的阶段重新定义标准比如“基部出现黄色区域”归为阶段 3而不是“黄色面积达到 20%”。另外在推理时如果模型对阶段 2 和阶段 3 的 softmax 分数非常接近我会输出 top-2 让产线工人确认而不是直接给一个可能错误的硬判。5.3 验证集准确率很高放到现场却翻车现象在工作室里验证集准确率 96%部署到产线边缘盒子后同一批照片准确率掉到 80% 出头而且错误集中在光线较暗的样本。原因训练集和验证集是从同一个视频里随机抽帧切分的。由于视频连续采集相邻帧很可能来自同一个菠萝的同一个成熟阶段随机切分导致验证集和训练集有些样本几乎是同一张图评估结果自然虚高。现场环境光线变了模型没有见过这种分布准确率就掉下来了。解决数据切分必须按拍摄批次或者菠萝个体编号划分绝不能在随机切分下直接上线。我会在收数据时记录每张图的拍摄时间和编号训练集用前三个批次的图验证集用后一个批次的图这样验证集才接近“没见过的菠萝”。还有一招是现场抽 500 张新图先不标注直接跑模型把 confident 置信度低的图挑出来让人工看能快速发现分布漂移。5.4 换用更大的输入分辨率准确率反而下降现象我觉得 224x224 比 160x160 能保留更多细节就把输入从 160 改到 224结果验证准确率没有提升反而低了 2 个点。原因成熟度分类真正有用的特征是颜色分布不是果眼纹理。分辨率上去了模型开始关注菠萝果眼的细节纹理而这些纹理在不同品种和个体之间差异很大属于噪声。ShuffleNet 是轻量网络容量有限过度学习这些纹理会牺牲对颜色整体分布的拟合。解决做一次输入分辨率的消融实验分别跑 160、192、224在同一个验证集上对比。我试过的菠萝项目里160 和 192 几乎打平224 反而差一些。最终我选了 192兼容大多数部署框架的算子优化。如果你还要识别菠萝表面的病虫害分辨率就需要更高那是另一个任务。5.5 预训练权重加载报尺寸不匹配现象运行代码时报size mismatch for fc.weight: copying a param with shape torch.Size([1000, 1024]) from checkpoint, the shape in current model is torch.Size([8, 1024])。原因这是把fc改成 8 类之后直接加载官方预训练权重老权重最后一个全连接层是 1000 类形状对不上。如果没用weightsDEFAULT而是手动下载了一个 state_dict最容易遇到这个问题。解决用strictFalse加载并丢掉分类头代码这样写import torch from torchvision.models import shufflenet_v2_x1_0 model shufflenet_v2_x1_0(weightsNone) model.fc nn.Linear(model.fc.in_features, 8) state_dict torch.load(shufflenetv2_x1_0.pth) state_dict.pop(fc.weight, None) state_dict.pop(fc.bias, None) model.load_state_dict(state_dict, strictFalse)strictFalse允许当前模型里某些层这里是新的 fc在 state_dict 中找不到对应权重。注意一定要先pop掉旧 fc 的权重否则 PyTorch 依然会报unexpected key。如果是不同 torchvision 版本之间切换可能出现层名前缀不同这时候可以先打印model.state_dict().keys()和state_dict.keys()用集合运算找出差异再决定改哪边。6. 用混淆矩阵和ONNX量化验收模型从训练到可以上线的最后一步模型训练完第一步不是部署而是把验证集上的混淆矩阵拉出来逐类看召回率。只看 overall acc 一定会漏掉问题比如阶段 1 全错被阶段 2 覆盖总准确率可能只少 3%但产线会丢失一批未成熟菠萝。用 sklearn 生成混淆矩阵from sklearn.metrics import confusion_matrix import numpy as np all_preds [] all_labels [] model.eval() with torch.no_grad(): for images, labels in val_loader: images images.to(device) preds model(images).argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) conf confusion_matrix(all_labels, all_preds) print(conf)读混淆矩阵时我习惯看每一行的对角线有没有形成“棒棒糖”形状。如果某个阶段的错误集中在相邻列说明模型学到了连续的颜色过渡如果错误跳到对角线远端那通常意味着标注有问题或者数据混入了异常图像。对菠萝项目来说相邻阶段互相错是常态只要每一行的大部分概率落在对角线或相邻格这个模型就有行业落地价值。验证通过后导出 ONNX 方便边缘部署。ShuffleNet 里的通道混洗在 ONNX 中会变成多个 transpose这是正常的但要在导出后验证输出是否一致model.eval() dummy_input torch.randn(1, 3, 192, 192) torch.onnx.export( model, dummy_input, shufflenet8.onnx, input_names[input], output_names[output], opset_version13, dynamic_axes{input: {0: batch}, output: {0: batch}} )参数说明opset_version13是我常用的最低版本太低会缺少一些 reshape 算子的支持dynamic_axes让批量维度动态化这样产线上一张一张推理或一次处理多张都行。导出后建议用 onnxruntime 加载跑一遍对比 PyTorch 输出的最大误差如果超过 1e-4就要检查是不是某些算子在 ONNX 里被错误融合。我最早做菠萝成熟度项目时没有做混淆矩阵就直接部署结果阶段 7 和阶段 8 被完全忽略因为它们在验证集里准确率很高但样本极少。后来每次训练结束都会先看混淆矩阵确认每一个成熟阶段都不缺失再进入部署流程。另一个习惯是把最有争议的阶段 4 到阶段 6 的误判样本打印成一张大图贴在产线旁边让工人知道模型在哪些地方会犯错比让工人盲目相信模型有价值得多。希望这套流程能帮你在自己的成熟度分类项目里少走几步弯路。本文还有配套的精品资源点击获取