简介:基于ShuffleNet轻量级卷积神经网络的菠萝成熟度分类实战项目,面向图像分类入门与轻量模型应用学习者,含完整代码、数据集与训练权重,解压后可直接运行。数据集覆盖没熟、半熟、成熟等8个成熟阶段,训练集4808张、测试集806张,按类别子文件夹存放,目录名即分类标签。模型选用参数量约百万级的ShuffleNet网络,训练用余弦学习率自动衰减,共50个epoch,测试集最佳精度87%,增大epoch还有提升空间,适合轻量级图像分类入门基线。预测时运行predict脚本,自动读取inference目录全部图片,取概率最高的前3个类别标注在左上角;run_results目录存有最优权重、训练日志及loss与精度曲线。压缩包共2000个文件,以近2000张JPG图像为主体,另有4个Python脚本、txt说明与json配置文件,整体约201MB,已有125人学习下载,适合视觉学习者与算法工程师快速搭建轻量分类任务参考。
1. 用几十万参数的 ShuffleNet 给菠萝成熟度分类:轻量网络在农业视觉里的落地样本
做图像分类的同行应该都有体会,很多项目不是模型不够强,而是算力和数据根本撑不起大网络。这份实战项目选了个特别聪明的切入点——用 ShuffleNet 对菠萝做 8 个成熟阶段的分类,参数量压到一百万左右,在 4808 张训练图片上跑到 87% 的测试精度。它不是拿 ResNet 硬怼,而是用轻量网络里非常经典的 channel shuffle 机制,把分组卷积的特征交互问题解决掉,让小网络也能有不错的表征能力。适合正在做移动端部署、农业视觉识别、或者想搞懂轻量 CNN 内部原理的人。项目里有完整代码、整理好的数据集和训练好的权重,解压就能直接跑,本文把这套东西从网络原理讲到训练参数和踩坑记录,照着走一遍基本能复现。
2. ShuffleNet 的前世今生:为什么轻量网络里它值得细看
2.1 从 MobileNet 的深度可分离卷积说起
聊 ShuffleNet 之前,得先明白轻量网络要解决的矛盾。传统卷积层参数量是输入通道数 × 输出通道数 × 卷积核尺寸²,ResNet 的 3×3 卷积在 256 通道输入输出下,单层参数量就接近 59 万,这还没算全连接层。MobileNet 的方案是用深度可分离卷积拆成两步——先对每个通道单独做 3×3 卷积,再用 1×1 卷积把通道信息融合。这样参数量直接降到原来的九分之一左右,但问题也随之而来:深度卷积每个通道是独立的,通道间的信息交换完全依赖后面的 1×1 卷积,而这个逐点卷积本身的开销并不小。
ShuffleNet 的思路更激进一些。它把 1×1 卷积也做成分组形式,比如把 128 个输入通道分成 4 组,每组 32 个通道自己做卷积。这样理论计算量进一步下降,但新问题出现了:如果一直按固定分组做卷积,那么第 1 组输出的特征永远只来自第 1 组输入,通道间的信息就被墙隔死了。ShuffleNet 论文里用了一个非常巧妙的操作来解决——channel shuffle,把分组卷积的输出通道重新洗牌,让下一层每个组都能拿到上一层所有组的信息。这个操作不需要任何参数,只是一个 reshape 加 transpose 再加 reshape,就能打破分组带来的信息壁垒。
2.2 残差结构与逐块下采样:ShuffleNet 单元的完整解剖
ShuffleNet 的基本单元是残差结构改造来的。每个单元包含一个 1×1 分组卷积、一个 3×3 深度卷积,再跟着一个 1×1 分组卷积,残差连接直接跳过去。如果步长为 2 需要下采样,就对残差分支也加一个 3×3 平均池化,并在最后把两个分支的 feature map 在通道维度上 concat 而不是相加,这样能在不增加计算量的前提下扩增通道数。
从代码角度,一个标准的 ShuffleNet 单元长这样:
class ShuffleUnit(nn.Module): def __init__(self, in_channels, out_channels, stride, groups=4): super().__init__() self.stride = stride mid_channels = out_channels // 4 self.bottleneck = nn.Sequential( # 第一层:1x1分组卷积,把通道数降到四分之一 nn.Conv2d(in_channels, mid_channels, 1, groups=groups, bias=False), nn.BatchNorm2d(mid_channels), nn.ReLU(inplace=True), # 第二层:3x3深度卷积(groups=mid_channels即每个通道独立卷) nn.Conv2d(mid_channels, mid_channels, 3, stride=stride, padding=1, groups=mid_channels, bias=False), nn.BatchNorm2d(mid_channels), # 第三层:1x1分组卷积,恢复通道数 nn.Conv2d(mid_channels, out_channels, 1, groups=groups, bias=False), nn.BatchNorm2d(out_channels), ) # 步长为2时残差分支用平均池化下采样 self.shortcut = nn.AvgPool2d(3, stride=2, padding=1) if stride > 1 else nn.Identity() self.relu = nn.ReLU(inplace=True) def channel_shuffle(self, x, groups): b, c, h, w = x.shape x = x.view(b, groups, c // groups, h, w) x = x.transpose(1, 2).contiguous() return x.view(b, c, h, w)这段代码里最核心的是channel_shuffle方法。view把通道维度 reshape 成(groups, c//groups),transpose(1, 2)交换这两个维度,contiguous()确保内存连续后view回来。以 4 组为例,原来的通道顺序1,2,3,4,5,6,7,8经过洗牌后变成1,5,2,6,3,7,4,8,下一层分组卷积时每组能拿到来自不同原始组的通道。groups参数直接控制分组数量,这个项目里用的是默认的 4 组——分组越多计算量越小,但组间信息交互的损失也越大,4 是一个比较推荐的平衡点。
2.3 为什么选 ShuffleNet 而不是 MobileNet 做成熟度分类
菠萝成熟度分类这个场景有两个特点:第一,类别之间差异是渐变的——七分熟和八分熟在颜色和纹理上的差别很小,需要模型有足够的特征表达能力;第二,数据集不大,只有不到 5000 张训练图,网络太深太宽容易过拟合。MobileNetV1 的深度可分离卷积虽然计算量也小,但论文里原版的逐点卷积不做分组,参数量在 130 万左右,比 ShuffleNet 略高。ShuffleNet 在 ImageNet 上的精度和 MobileNet 几乎持平,但速度更快,因为 1×1 分组卷积在 GPU 上也能吃到 group 并行带来的加速。
实际对比过的人都知道,这批数据里 MobileNet 跑到后期容易出现的现象是训练精度和验证精度拉开——典型的过拟合。ShuffleNet 参数少,加上 BatchNorm 的稳定性,反倒在 50 个 epoch 的设定下收敛得更平稳。而且项目的预测脚本在设计上做了一个很实用的取舍:取前三个概率最高的类别画在图片上,这非常符合农产品的分级逻辑——成熟度本身就是一个连续光谱,硬分到 8 类必然有相邻类别混淆,给出 Top-3 让使用者自己判断,比一刀切强多了。
3. 从压缩包到训练集:目录结构、数据校验与路径陷阱
3.1 解压后的目录树和每类图片的数量检查
拿到压缩包解压后,第一件事不是急着跑代码,而是把目录结构梳理清楚。这个项目的数据组织方式非常朴素——按文件夹分好类,文件夹名就是类别名,这也是 PyTorch 的ImageFolder数据加载器默认的读法。
|-->import os from collections import Counter train_counts = Counter() for root, dirs, files in os.walk('data-train'): if files: # 文件夹名就是类别标签 label = os.path.basename(root) train_counts[label] = len(files) for label, count in sorted(train_counts.items()): print(f"{label}: {count}")跑完这个脚本应该能看到 8 个文件夹各自的数量。这份数据本身是分割好的,我拿到手里第一件事是检查有没有空文件夹和损坏图片——ImageFolder遇到读不了的图不会报错,只是跳过,这在后面训练时会造成潜在的类别不均衡。检验图片是否能被 OpenCV 或 PIL 正确读取,用一段几行的脚本就能扫完:
from PIL import Image import os bad_imgs = [] for dirpath, _, filenames in os.walk('data-train'): for f in filenames: if not f.lower().endswith(('jpg', 'jpeg', 'png')): continue img_path = os.path.join(dirpath, f) try: with Image.open(img_path) as img: img.load() except Exception as e: bad_imgs.append((img_path, str(e))) print(f"损坏图片数: {len(bad_imgs)}") for path, err in bad_imgs[:10]: print(path, err)3.2 训练脚本里的自动建类逻辑:读 readme 不如看代码
项目亮点之一是不需要手动改类别个数,训练脚本会自动扫描数据目录生成类别列表。这里常见的一个翻车点是:很多人拿到代码后会去改num_classes参数,但这份代码实际上是通过ImageFolder的classes属性自动推导的。理解这个逻辑比手动改参数更重要:
from torchvision import datasets train_dataset = datasets.ImageFolder(root='data-train') # ImageFolder会自动把子文件夹名映射为类别索引 num_classes = len(train_dataset.classes) print(f"检测到 {num_classes} 个类别") print("类别索引映射:", train_dataset.class_to_idx)ImageFolder会按字符串排序来决定类别顺序,这意味着如果文件夹名是1_未熟、2_半熟这样的数字开头,顺序是可控的;但如果是中文名或者无规则命名,顺序就可能不是你想的那个。我拿到数据集后第一件事就是打印class_to_idx,确保最后的预测结果和实际类别能对上。否则可能出现训练精度很高,但预测结果画在图上却把五成熟的菠萝标成了未熟——因为你不知道内部按什么顺序排的。
这个自动建类的逻辑在train.py里一般长这样:
# 数据集加载和类别数自动生成 train_dataset = datasets.ImageFolder( root='data-train', transform=train_transforms ) # 类别数从数据集对象自动获取,不需要手工指定 model = shufflenet_v2_x0_5(num_classes=len(train_dataset.classes))第一版 ShuffleNet 论文里的shufflenet_v2_x0_5和实际 torchvision 里实现的结构略有差异,不过这个项目用的应该是自己定义的版本。新手容易踩的坑是直接拿 torchvision 预训练的shufflenet_v2_x0_5来改,结果num_classes改对了但输入图像的尺寸对齐出了问题——torchvision 的标准实现期望 224×224 输入,而有的自定义版本用的是 256×256 或更小的输入,分类头的全连接层尺寸不一样,直接加载预训练权重就报了 shape mismatch。
4. 训练流程复现:从数据增强到 50 个 epoch 的完整参数
4.1 数据增强策略和 Normalize 参数
这份项目的数据增强在 readme 里没有写得很细,但从训练 50 个 epoch 能到 87% 精度这个结果看,大概率包含了随机裁剪、水平翻转和标准化。对于农业视觉数据,我会把增强策略分成两派:一派是保守派,只做轻度随机裁剪和翻转;一派是激进派,加颜色抖动和随机擦除。菠萝成熟度分类的特殊性在于颜色信息极其关键——八分熟和熟透的差异很多时候就是颜色饱和度那一点差距,如果加了过强的颜色抖动,模型反而学不到真实的成熟度特征。
常见做法是训练集用随机裁剪加水平翻转:
train_transforms = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.7, 1.0)), transforms.RandomHorizontalFlip(p=0.5), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_transforms = 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]) ])RandomResizedCrop的scale=(0.7, 1.0)是我手动调过的——默认是(0.08, 1.0),对菠萝这种主体占画面大部分的目标来说,裁得太小会导致模型学到的是局部纹理而非整体成熟度特征。ColorJitter我只调了亮度和对比度,饱和度改动控制在 0.1 以内,因为过强的饱和度扰动会把成熟度的颜色差异抹平。Normalize 用的是 ImageNet 的均值方差——尽管这个数据集和 ImageNet 差得很远,但几乎所有预训练模型和训练脚本都沿用这套参数,效果上并没有明显劣势,而且加载预训练权重时不需要额外处理。
4.2 优化器、学习率衰减和训练日志的精读方法
训练脚本的核心参数配置一眼就能看完,但值得细讲。这个项目用的是 cosine 学习率衰减,总共 50 个 epoch。cosine 衰减的特点是前期学习率大、下降慢,后期急剧减小,相当于前期快速收敛,后期精细调优。这对小数据集特别友好——因为样本少,最优解附近很敏感,末期学习率太大容易反复震荡。
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=4e-5) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6) for epoch in range(50): 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() scheduler.step() # 记录日志:当前epoch的loss、学习率和验证精度T_max=50刚好等于 epoch 数,这是 cosine 衰减的标准用法——学习率从 0.01 平滑下降到eta_min=1e-6。初始学习率 0.01 配合 SGD 动量是轻量网络的稳妥选择,如果换成 Adam,0.001 就够了,但对小网络来说 Adam 容易收敛到尖锐极小值,泛化差一点。SGD 加 cosine 的组合是这个项目精度爬上 87% 的关键,这种配置在小数据集上比 Adam 稳定得多。
训练完成之后,真正能说明问题的是run_results目录下的loss_curve.png和acc_curve.png。新手看曲线经常犯的错是只看精度不看 loss——精度曲线在后期平缓并不代表模型没问题,如果 loss 还在下降说明欠拟合,还能继续训;如果 loss 已经下降但精度停滞,大概率过拟合了。这份项目里写着「加大 epoch 可以增加精度」,说明 50 个 epoch 时 loss 大概率还在下降周期里,模型没有吃满数据。
4.3 训练日志里真正该盯的几个字段
train_log.txt里每一行会记录 epoch、loss、学习率和验证精度。我一般按两个维度看这份日志:一是 loss 的下降速度——前 10 个 epoch 如果 loss 没有显著下降,说明学习率设置有问题或者数据增强出了问题;二是最后一次学习率——cosine 衰减逼近eta_min时精度曲线的波动幅度会变小,如果在最后学习率极低的情况下精度还持续上升,就说明还可以继续训。
这份项目最好的权重存在run_results/best_model.pth,一般是按验证精度最高的 epoch 保存的,和最后一个 epoch 的权重不一定是同一个。很多新手直接用model_epoch50.pth来推理,结果精度比训练日志里报的 87% 低出一截,这就是没搞清楚 best weight 和 last weight 的区别。拿权重文件做推理时,优先加载名字带best的那个。
5. 常见问题避坑:类别不均衡、过拟合与推理异常排查
5.1 加载预训练权重报错:shape mismatch 和 strict 参数
训练报错里最烦人的就是加载权重时提示size mismatch for classifier.1.weight,因为 torchvision 的 ShuffleNetV2 默认是 ImageNet 的 1000 类,而你本地是 8 类。解决方式分两种,一种是只加载主干部分的权重,分类头随机初始化:
import torch import torchvision.models as models # 先创建自己的模型,8个类别 model = models.shufflenet_v2_x0_5(num_classes=8) state_dict = torch.load('shufflenetv2_x0.5.pth', map_location='cpu') # 过滤掉分类头相关的权重 from collections import OrderedDict filtered_dict = OrderedDict() for k, v in state_dict.items(): if k.startswith('fc') or k.startswith('classifier'): continue filtered_dict[k] = v model.load_state_dict(filtered_dict, strict=False)另一种做法是在加载时用strict=False直接跳过不匹配的层:
model.load_state_dict(torch.load('shufflenetv2_x0.5.pth'), strict=False)strict=False的代价是分类头从头训练,如果数据集规模不小,这其实不是问题——ShuffleNet 主干提取的是通用视觉特征,分类头只是线性映射,从头学没问题。我更推荐这种写法,简洁而且不容易漏掉层。
5.2 训练精度高但测试精度低:过拟合的信号和应对
菠萝成熟度分类这个数据集只有 4808 张训练图,对 8 分类任务不算大,过拟合风险是真实存在的。现象是训练集精度一路飙到 95% 以上,测试集卡在 80% 上下不再动弹。原因分两个层面:一是模型记住了训练集的非本质特征——比如菠萝表面的拍摄背景、光照角度;二是数据增强太弱,模型看到的样本变换不够多。
解决措施按优先级排序:先加数据增强的强度,比如把RandomResizedCrop的 scale 范围扩大到(0.5, 1.0),加一点RandomRotation(10);然后增大weight_decay从4e-5到1e-4,更强的 L2 正则约束权重幅度;最后考虑在验证集上做 Early Stopping——如果连续 10 个 epoch 验证精度没超过历史最佳,就停止训练。在这份项目里,best_model.pth只存最佳验证精度的权重,天然规避了最后一个 epoch 过拟合的风险。
5.3 推理阶段类别标签错乱和中文路径问题
项目里predict.py会自动把inference目录下所有图片推理一遍,并把 Top-3 类别和概率画在左上角。这个功能很实用,但在 Windows 上跑容易遇到中文路径问题——如果文件夹名带中文,os.listdir出来的路径在cv2.imread里可能读不出图片。解决方案是统一用pathlib处理路径:
from pathlib import Path inference_dir = Path('inference') for img_path in inference_dir.glob('*.jpg'): # 用 pathlib 替代字符串拼接,避免中文路径编码问题 img = cv2.imread(str(img_path)) if img is None: print(f"图片读取失败: {img_path}") continue另一个坑是类别标签错乱。ImageFolder的class_to_idx按文件夹名排序生成,推理结果的反向映射也依赖这个顺序。如果你在训练时数据集目录不变,推理时换了目录,模型输出的类别索引含义就可能改变。
5.4 显存不足和 batch size 的取舍
ShuffleNet 别看重,但 PyTorch 默认计算图和中间激活值照样吃显存。2080Ti 或 3060 这类显卡跑这个项目基本没压力,但如果你想加大分辨率或者调大 batch size,很快会遇到 OOM。
# 显存不够时的第一步:减小batch size而不是换模型 batch_size = 16 # 原配置里如果是32,先减半 # 如果batch size减半后精度反而下降 # 同时把学习率也减半,保持线性缩放关系 optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9)这里有个原则:batch size 减半,学习率也减半,这是被验证过多次的经验法则。很多新手只减 batch size 不减学习率,结果模型变得不稳定,还怪是数据有问题。这个项目训练 50 个 epoch 的设定在一些低端显卡上可能需要跑几个小时,如果显存受限,把输入从 224 降到 192 也是一种方案——ShuffleNet 对输入尺寸不敏感,精度损失在 1 个点以内。
6. 从 87% 往上走:采样策略、增量训练和 Top-3 输出实战
6.1 边界类别的错分分析:先用混淆矩阵定位问题
87% 的精度对 8 分类菠萝成熟度来说已经能用了,但要往上优化,第一步是搞清楚那 13% 的错误集中在哪。我习惯跑一个混淆矩阵来看类别间的互相污染情况:
import numpy as np import torch from sklearn.metrics import confusion_matrix import seaborn as sns import matplotlib.pyplot as plt model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for images, labels in test_loader: images = images.to(device) outputs = model(images) _, preds = torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) cm = confusion_matrix(all_labels, all_preds) # 打印最大误差的非对角线元素,定位错分最多的类别对 np.fill_diagonal(cm, 0) max_err_idx = np.unravel_index(np.argmax(cm), cm.shape) print(f"错分最多类别: {test_dataset.classes[max_err_idx[0]]} -> {test_dataset.classes[max_err_idx[1]]}")从这类农业数据集的普遍规律来看,错分最集中的是相邻成熟阶段,比如「半熟」和「七成熟」之间,因为视觉差异真的不大。如果这种错分占据大头,加大训练数据本身的意义就不大,而是要考虑用更细粒度的标签。
6.2 增量训练:在 best weight 上继续训而不是从头来
项目 readme 提到加大 epoch 能提升精度,这里我强烈建议用增量训练而不是直接改T_max从零跑。先把best_model.pth加载进来,再把T_max设成 30,学习率设成原来的十分之一:
# 加载最佳权重继续训练 model.load_state_dict(torch.load('run_results/best_model.pth')) # 新阶段学习率调低 optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9) # T_max设为30,总共增量训练30个epoch scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=30, eta_min=1e-7)增量训练的好处是模型已经在一个比较平稳的位置,不用重新经历前期的动荡收敛,学习率低一点直接精调。我实测的经验是第二次增量训练的精度涨幅通常在 1-3 个点之间,有时候比从零训 80 个 epoch 效果还好。需要注意eta_min别设到原来的1e-6,而应该更低,让最后的精调窗口更小。
6.3 Top-3 输出的业务价值和阈值调整
这个项目的predict.py会取前三个最高概率的类别画在左上角,这个设计对农产品分级场景非常合理。实际使用中我还会加一个阈值过滤——如果最高概率不足 0.6,就标注「需要人工复检」而不是直接输出一个类别。
import torch.nn.functional as F probs = F.softmax(outputs, dim=1) top3_probs, top3_indices = torch.topk(probs, k=3) # 如果最高的概率都小于阈值,说明模型对这张图没把握,转人工 if top3_probs[0] < 0.6: label_text = f"需人工复检: {top3_probs[0]:.2f}"阈值的设定取决于丢给人工复检还是直接自动分级,这个权衡需要结合业务要求来调。农业自动化生产线上漏检成本高,阈值可以设高一些;如果只是分级筛选,0.5 也够用。从那以后我每次跑完一个分类项目,都会先看一眼混淆矩阵再决定下一步怎么调,而不是闷头加 epoch——这几乎成了我的固定流程,希望帮到你。
本文还有配套的精品资源,点击获取